2923. Tree

The hanging tree contains n (1 ≤ n ≤ 10⁶) vertices. Each vertex is colored in one of n colors. For each vertex v find the number of different colors, occurring in the subtree with root v.

Input. The first line contains the number n. The next n lines describe the vertices. The description of the vertex i has the form p_i c_i, where p_i is the parent number of the vertex i, and c_i is the color of the vertex i (1 ≤ c_i ≤ n). For the root of the tree p_i = 0.

Output. Print n numbers, denoting the number of different colors in the subtrees rooted at the vertices 1, 2, ..., n.

Sample input	Sample output
5 2 1 3 2 0 3 3 3 2 1	1 2 3 1 1

SOLUTION

graphs, depth first search

Algorithm analysis

Let’s start a depth-first search from the root of the tree. For each vertex i, store a set s_i, where we’ll accumulate the colors of the vertices in its subtrees. If j is the son of vertex i during a depth-first search, then s_j must be included in s_i. The number of distinct colors in the subtree rooted at i equals the size of the set s_i.

Example

On the left, near each vertex, its color is given. On the right, near each vertex, the set of colors in a subtree rooted at it is given.

Algorithm realization

In color[i] keep the color of the i-th vertex. In the set s[i] we’ll accumulate colors in the subtree with root i. Store the directed tree in the adjacency list g. In res[i] store the number of different colors in the subtree rooted at i.

#define MAX 1000010

int color[MAX], res[MAX];

set<int> s[MAX];

vector<vector<int> > g;

Function dfs starts the depth-first search from the vertex v. Initially, put in s[v] the color of the vertex v. For each edge of the tree (v, to) add the set s[to] to s[v]. The number of different colors in the subtree with the root v equals the size of the set s[v], store it in res[v].

void dfs(int v)

{

int i, to;

s[v].insert(color[v]);

for(i = 0; i < g[v].size(); i++)

{

to = g[v][i];

dfs(to);

If the size of the set s[v] is less than the size of the set s[to], swap them. Next, the content of the smaller set s[to] is added to the set s[v].

if (s[v].size() < s[to].size()) s[v].swap(s[to]);

s[v].insert(s[to].begin(), s[to].end());

Clear the set s[to] – it will no longer be useful to us.

s[to].clear();

}

res[v] = s[v].size();

}

The main part of the program. Read the input data.

scanf("%d",&n);

g.resize(n+1);

for(i = 1; i <= n; i++)

{

scanf("%d %d",&p,&c);

g[p].push_back(i);

color[i] = c;

}

Start the depth-first search from the root of the tree – the vertex number zero.

dfs(0);

Print the answer.

for(i = 1; i <= n; i++)

printf("%d ",res[i]);

printf("\n");