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.
classDisjointSet:def__init__(self,nodes):self.parent={x:xforxinnodes}self.rank={x:0forxinnodes}deffind(self,x):ifself.parent[x]!=x:self.parent[x]=self.find(self.parent[x])returnself.parent[x]defadd_node(self,x):ifxnotinself.parent:self.parent[x]=xself.rank[x]=0defunion(self,x,y):self.add_node(x)self.add_node(y)x_root=self.find(x)y_root=self.find(y)ifx_root==y_root:returnifself.rank[x_root]<self.rank[y_root]:x_root,y_root=y_root,x_rootself.parent[y_root]=x_rootself.rank[x_root]+=self.rank[y_root]defsolve(edges):nodes=set()foru,v,colorinedges:nodes.add(u)nodes.add(v)ds=DisjointSet(nodes)foru,v,colorinedges:# Two nodes can be connected through a black edge. Any other node connected through a black edge to these previous nodes will then belong to the same set of nodes all connected through black edges. A node that belongs to a black edge subset cannot be connected with another node belonging to the same or another black edge subset.ifcolor=='b':ds.union(u,v)subsets_sizes={}fornodeinnodes:root=ds.find(node)subsets_sizes[root]=subsets_sizes.get(root,0)+1n=len(nodes)total_triplets=n*(n-1)*(n-2)// 6invalid_triplets=0forsizeinsubsets_sizes.values():ifsize>=3:invalid_triplets+=size*(size-1)*(size-2)// 6ifsize>=2:invalid_triplets+=size*(size-1)// 2 * (n - size)valid_triplets=total_triplets-invalid_tripletsreturnvalid_triplets%(10**9+7)n=int(input().strip())edges=[]for_inrange(n-1):u,v,color=input().strip().split()edges.append((int(u),int(v),color))print(solve(edges))
Cookie support is required to access HackerRank
Seems like cookies are disabled on this browser, please enable them to open this website
Kundu and Tree
You are viewing a single comment's thread. Return to all comments →