We use cookies to ensure you have the best browsing experience on our website. Please read our cookie policy for more information about how we use cookies.
Here's my O(n) solution in Java.
It works as follows:
It constructs disjoint components consisting of vertices only connected by black edges using disjoint sets (i.e. Union-Find sets)
It calculates the number of all possible triplets, i.e. binom(n, 3): "from n elements choose 3 without regarding the order"
For each distinct component C_i subtract from the number of possible triplets:
the number of triplets within the component, i.e. binom(n_i, 3) where n_i is the size of component C_i
the number of triplets build from 2 vertices of the component and 1 other (not in the component), i.e. binom(n_i, 2) * (n - n_i)
Print the result modulo 10^9 + 7
importjava.io.*;importjava.util.*;importjava.text.*;importjava.math.*;importjava.util.regex.*;classDisjointSet{DisjointSetparent=this;intsize=1;DisjointSetfindRoot(){if(parent!=this){// path compressionparent=parent.findRoot();}returnparent;}voidunion(DisjointSetother){if(other==this)return;DisjointSetroot=findRoot();DisjointSetotherRoot=other.findRoot();if(otherRoot==root)return;// union by sizeif(root.size>=otherRoot.size){otherRoot.parent=root;root.size+=otherRoot.size;}else{root.parent=otherRoot;otherRoot.size+=root.size;}}}publicclassSolution{finalstaticlongmod=1_000_000_007;publicstaticvoidmain(String[]args){Scannerscanner=newScanner(System.in);intn=scanner.nextInt();DisjointSet[]components=newDisjointSet[n+1];for(inti=0;i<n-1;i++){inta=scanner.nextInt();intb=scanner.nextInt();charcolor=scanner.next().charAt(0);// ignore red edgesif(color=='r')continue;DisjointSetaComponent=createComponentIfNeeded(components,a);DisjointSetbComponent=createComponentIfNeeded(components,b);aComponent.union(bComponent);}// uniquify componentsSet<DisjointSet>uniqueComponents=newHashSet<>();for(inti=0;i<n;i++){DisjointSetcomponent=components[i];if(component!=null){uniqueComponents.add(component.findRoot());}}longvalidTriplets=choose3from(n);// all possible tripletsfor(DisjointSetcomponent:uniqueComponents){// subtract all triplets within the componentsvalidTriplets-=choose3from(component.size);// subtract all triplets build from 2 vertices of the components and 1 other vertexvalidTriplets-=choose2from(component.size)*(n-component.size);}// apply modulo for some reason beyond my imaginationSystem.out.println(validTriplets%mod);}// lazily create a componentstaticDisjointSetcreateComponentIfNeeded(DisjointSet[]components,intindex){if(components[index]==null){components[index]=newDisjointSet();}returncomponents[index];}// optimized calculation of binom(n, 3)staticlongchoose3from(intn){if(n<3)return0;// k = 3// n! / ( k! (n-k)! ) = (n-k+1) * ... * n / k! = (n - 2) * ... * n / 6longres=1;for(intx=n-2;x<=n;x++){res*=x;}returnres/6;}// optimized calculation of binom(n, 2)staticlongchoose2from(intn){if(n<2)return0;// k = 2// n! / (k! (n-k)!) = (n-k+1) * ... * n / k! = (n - 1) * ... * n / 2longres=1;for(intx=n-1;x<=n;x++){res*=x;}returnres/2;}}
Kundu and Tree
You are viewing a single comment's thread. Return to all comments →
Here's my O(n) solution in Java. It works as follows: