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.
importsys,math# Pretty important line here.sys.setrecursionlimit(10**6)defpreprocess_original_tree(node,parent):# "Jumping" in the current node.globaltimertin[node]=timertimer+=1# Building a chart of the 2^i ancestors of each node.up[node][0]=parent# Setting each node's (which is the current child) 2^i ancestors' value.foriinrange(1,LOG):ifup[node][i-1]!=-1:up[node][i]=up[up[node][i-1]][i-1]# Perform DFS to set current node's children's depth recursively.forneighborinadj[node]:ifneighbor==parent:continuedepth[neighbor]=depth[node]+1preprocess_original_tree(neighbor,node)# "Jumping" out of the current node.tout[node]=timertimer+=1deflift_node(node,k):# Jumping to the greatest 2^i ancestor each time.foriinrange(LOG-1,-1,-1):ifk&(1<<i):node=up[node][i]returnnodedefget_lca(u,v):# Equalizing both node's depths.ifdepth[u]<depth[v]:u,v=v,uu=lift_node(u,depth[u]-depth[v])ifu==v:returnu# Jumping to the greatest 2^i ancestor each time.foriinrange(LOG-1,-1,-1):ifup[u][i]!=up[v][i]:u=up[u][i]v=up[v][i]returnup[u][0]defget_distance(u,v):ancestor=get_lca(u,v)# It uses the original tree's preprocessed depths.returndepth[u]+depth[v]-2*depth[ancestor]defbuild_virtual_tree(nodes):# Adding relevant nodes to virtual tree.nodes.sort(key=lambdax:tin[x])m=len(nodes)vt_nodes=nodes[:]foriinrange(m-1):vt_nodes.append(get_lca(nodes[i],nodes[i+1]))vt_nodes=list(set(vt_nodes))vt_nodes.sort(key=lambdax:tin[x])# Connecting nodes in virtual tree.tree={node:[]fornodeinvt_nodes}stack=[]# All virtual tree nodes are stored in the order in which they were found, thus preserving their hierarchy from left to right.fornodeinvt_nodes:# Validating if the last ancestor in the stack is the ancestor of the current node.whilestackandtout[stack[-1]]<tin[node]:stack.pop()ifstack:tree[stack[-1]].append(node)stack.append(node)returntree,vt_nodesdefsolve_query(query_nodes):# Traversing query nodes (virtual tree's nodes)defdp(u):nonlocalress=query_val.get(u,0)# Performing DFS.forvinvt[u]:sub=dp(v)# Since # sum(u in sub) * sum(v not in sub) # = (sum(u in sub)) * (sum(v not in sub)) # = sub * (S_total - sub)res=(res+sub*(S_total-sub)%MOD*get_distance(u,v))%MODs+=sub# Returning the total sum of the current subtree.returnsiflen(query_nodes)<2:return0S_total=sum(query_nodes)query_val={node:nodefornodeinquery_nodes}vt,vt_nodes=build_virtual_tree(query_nodes)res=0dp(vt_nodes[0])returnresMOD=10**9+7timer=0data=sys.stdin.read().split()it=iter(data)n=int(next(it))q=int(next(it))LOG=int(math.log2(n))+1up=[[-1]*LOGfor_inrange(n+1)]depth=[0]*(n+1)tin=[0]*(n+1)tout=[0]*(n+1)# Building original tree.adj=[[]for_inrange(n+1)]for_inrange(n-1):u,v=int(next(it)),int(next(it))adj[u].append(v)adj[v].append(u)preprocess_original_tree(1,-1)res=[]for_inrange(q):k=int(next(it))query_nodes=[int(next(it))for_inrange(k)]res.append(str(solve_query(query_nodes)))sys.stdout.write("\n".join(res))
Cookie support is required to access HackerRank
Seems like cookies are disabled on this browser, please enable them to open this website
Kitty's Calculations on a Tree
You are viewing a single comment's thread. Return to all comments →