Coloring Tree

                            Solution in C :

In   C++  :







#include<cstdio>
#include<cstring>
#include<queue>
#include<vector>
#include<map>
#include<iostream>
#include<string>
#include<algorithm>

using namespace std;

#define foreach(e, x) for(__typeof(x.begin()) e = x.begin(); e != x.end(); ++ e)
#define next nxt

const int N = 100000 + 10;
const int MOD = 1000000000 + 7;

int n, m, r;
int tot;
vector<int> adj[N];
int color[N];
int father[N];
int sqn[N];
int ret[N];
int ord[N][2];
int next[N];
int c[N];
pair< pair<int, int>, int> query[N];

void dfs(int u)
{
	sqn[tot] = color[u];
	ord[u][0] = tot ++;
	foreach(it, adj[u]) {
		int v = *it;
		if (v == father[u]) continue;
		father[v] = u;
		dfs(v);
	}
	ord[u][1] = tot;
}

void add(int u, int x)
{
	for( ; u <= n; u += u & -u)
		c[u] += x;
}

int ask(int u)
{
	int ret = 0;
	for( ; u; u -= u & -u)
		ret += c[u];
	return ret;
}

void solve()
{
	cin >> n >> m >> r;
	int u, v;
	for(int i = 1; i < n; ++ i) {
		scanf("%d%d", &u, &v);
		-- u, -- v;
		adj[u].push_back(v);
		adj[v].push_back(u);
	}
	-- r;
	father[r] = -1;
	for(int i = 0; i < n; ++ i) {
		scanf("%d", &color[i]);
	}
	tot = 0;
	dfs(r);
	for(int i = 0; i < n; ++ i) {
		query[i] = make_pair(make_pair(ord[i][0], ord[i][1]), i);
	}
	sort(query, query + n);
	map<int, int> s;
	for(int i = n - 1; i >= 0; -- i) {
		if (s.count(sqn[i]) == 0) {
			next[i] = -1;
		} else {
			next[i] = s[sqn[i]];
		}
		s[sqn[i]] = i;
	}
	foreach(it, s) {
		add(it->second + 1, 1);
	}
	int l, r;
	int ptr = 0;
	for(int i = 0; i < n; ++ i) {
		l = query[i].first.first;
		r = query[i].first.second;
		for ( ; ptr < l; ) {
			if (next[ptr] != -1)
				add(next[ptr] + 1, 1);
			++ ptr;
		}
		ret[query[i].second] = ask(r) - ask(l);
	}
	for(int i = 0; i < n; ++ i) {
	}
	for(int i = 0; i < m; ++ i) {
		scanf("%d", &u);
		-- u;
		printf("%d\n", ret[u]);
	}
}

int main()
{
	solve();
	return 0;
}








In   Java  :






import java.io.*;
import java.util.*;
import java.text.*;
import java.math.*;
import java.util.regex.*;

public class ColoringTree {
public static class TreeNode {
ArrayList<TreeNode> ch;
int color, numberOfColors;
boolean visited;

TreeNode() {
ch = new ArrayList<TreeNode>();
visited = false;
}
}

public static TreeSet<Integer> colors(TreeNode node) {
TreeSet<Integer> combined;
node.visited = true;
combined = new TreeSet<Integer>();
LinkedList<TreeSet<Integer>> childColors 
= new LinkedList<TreeSet<Integer>>();
for (TreeNode child : node.ch) {
if (child.visited)
continue;
TreeSet<Integer> col = colors(child);
if (col.size() > combined.size())
combined = col;
childColors.add(col);
}
for (TreeSet<Integer> col : childColors) {
if (col != combined)
combined.addAll(col);
}
combined.add(node.color);
node.numberOfColors = combined.size();
return combined;
}

public static void main(String[] args) {
int n, m, root;
Scanner in = new Scanner(System.in);
n = in.nextInt();
m = in.nextInt();
root = in.nextInt();
TreeNode[] nodes = new TreeNode[n + 1];
for (int i = 1; i <= n; i++)
nodes[i] = new TreeNode();
for (int i = 0; i < n - 1; i++) {
int a, b;
a = in.nextInt();
b = in.nextInt();
nodes[a].ch.add(nodes[b]);
nodes[b].ch.add(nodes[a]);
}
for (int i = 1; i <= n; i++) {
nodes[i].color = in.nextInt();
}
colors(nodes[root]);
for (int i = 0; i < m; i++) {
System.out.println(nodes[in.nextInt()].numberOfColors);
}

}
}








In   Python3 :







from collections import Counter
n, m, root = map(int, input().split())
uniquenum = dict()
multipleset = dict()
adj = dict()
for _ in range(n-1):
    n1, n2 = map(int, input().split())
    if n1 in adj:
        adj[n1].add(n2)
    else:
        adj[n1] = set([n2])
    if n2 in adj:
        adj[n2].add(n1)
    else:
        adj[n2] = set([n1])

colors = [int(input()) for _ in range(n)]
multiples = set(Counter(colors)-Counter(set(colors)))
colors.insert(0, 0)
totalcolors = len(set(colors[1:]))

stack = [root]
added = set([root])
visited = set()
while len(stack)>0:
    node = stack[len(stack)-1]
    if node not in visited:
        visited.add(node)
        for child in adj[node]-added:
            stack.append(child)
            added.add(child)
    else:
        if colors[node] in multiples:
            uniquenum[node] = 0
            multipleset[node] = set([colors[node]])
        else:
            uniquenum[node] = 1
            multipleset[node] = set()
        for child in adj[node]-added:
            uniquenum[node] += uniquenum[child]
            multipleset[node] |= multipleset[child]
        stack.pop()
        added.remove(node)   

for _ in range(m):
    node = int(input())
    print(uniquenum[node]+len(multipleset[node]))