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
An unexpected error occurred. Please try reloading the page. If problem persists, please contact support@hackerrank.com
Kundu and Tree
You are viewing a single comment's thread. Return to all comments →