HackerRank/Algorithms/Graph Theory/The Story of a Tree

Problem Summary

Alice and Bob play a game.

Given a tree with N nodes, two integers G and K, and G guesses (u1,v1),…,(ug,vg). (1<=N,G,K<=100000)

(ui,vi) means Alice guesses that ui is the parent of vi.

When a node of the tree is randomly selected (with equal probability) to be root by Bob, and Alice does not know, find the probability that the number of right guesses is greater than or equal to K.

Solution

At first I think of a solution with time complexity O(NG).

Let w[i] be the number of right guesses when node i is the root Bob picked. Then we only need to find w[i] of every node and count the ones greater than or equal to K.

Step 1. First, I pick node 1 as the root, and use DFS to find the parent of each node.

Step 2. Then, for each guess (ui,vi).

(1) If parent[vi]==ui, i.e. the guess is right when the root is node 1, or any other node that is on the same “side” with node 1, of edge (ui,vi). Then we add 1 to w[1] and w of the nodes that satisfying the mentioned condition.

(2) If parent[vi]!=ui. Since it is guaranteed that an undirected edge connecting ui and vi exists in the tree, we have parent[ui]==vi. Then we add 1 to w[ui] and w of the nodes on the same “side” with node ui of edge (ui,vi).

Step 3. Finally, we can check the w value of all nodes.

The time complexity is O(N*G), too high for N=G=100000. So we need some optimization.

Let us look back on the two kinds of nodes in step 2. The nodes in 2.(1) are the nodes that are not in the subtree with node vi as root. And the nodes in 2.(2) are the nodes of the subtree with node ui as root. This lead me to think about segment tree, which is much faster when it comes to updating and query. But we can not apply segment tree directly here. Instead, we can “split” w values of each node to itself and its ancestors. Let us assume that the REAL w value of node i equals the sum of its w value and w values of all ancestors of node i.

Then, when we want to add 1 to all w values of a subtree, we only need to add 1 to the w value of its root. So step 2 and 3 is changed as follows.

Step 2. For each guess (ui,vi),

(1) If parent[vi]==ui, add 1 to w1 and subtract 1 from w[vi].
(2) If parent[ui]==vi, add 1 to w[ui].

Step 3. Use DFS to find the REAL w value of each node, and find the answer.

In this way, we reduce the time complexity to O(N+G).

Code

#include <cmath>
#include <cstdio>
#include <cstring>
#include <iostream>

using namespace std;
const int maxn = 100002;
const int root = 1;

int n,ans,kk;
int tot,head[maxn<<1],next_e[maxn<<1],v[maxn<<1];

int parent[maxn],w[maxn];

void add(int a,int b)
{
	tot++; next_e[tot]=head[a]; head[a]=tot; v[tot]=b;
}

int gcd(int a,int b)
{
	if(a%b==0) return b;
	return gcd(b,a%b);
}

void dfs(int x,int p)
{
	parent[x]=p;
	for(int i=head[x];i!=0;i=next_e[i])
	{
		int y=v[i];
		if(parent[y]==0)
			dfs(y,x);
	}
}

void calc(int x,int sum)
{
	sum+=w[x];
	if(sum>=kk) ans++;

	for(int i=head[x];i!=0;i=next_e[i])
	{
		if(v[i]!=parent[x])
			calc(v[i],sum);
	}
}

int main() {

	freopen("a.in","r",stdin);

	int T,g;
	scanf("%d",&T);
while(T)
{
	T--;

	tot=0;
	memset(head,0,sizeof(head));
	memset(next_e,0,sizeof(next_e));

	int a,b;
	scanf("%d",&n);
	for(int i=1;i<n;i++)
	{
		scanf("%d%d",&a,&b);
		add(a,b);
		add(b,a);
	}

	memset(parent,0,sizeof(parent));
	dfs(root,-1);
	
	memset(w,0,sizeof(w));
	scanf("%d%d",&g,&kk);
	for(int i=1;i<=g;i++)
	{
		scanf("%d%d",&a,&b);
		if(parent[b]==a)
		{
			w[root]++;
			w[b]--;
		}
		else
		{
			/* It's guaranteed that an undirected edge 
			connecting a and b exists in the tree.
			So we have parent[a]==b */
			w[a]++;
		}
	}

	ans=0;
	calc(root,0);

	if(ans>0)
	{
		int c = gcd(ans,n);
		ans/=c;
		n/=c;
	}
	else
		n=1;

	printf("%d/%d\n",ans,n);
}
    return 0;
}