Given a rooted tree T, with ‘n’ nodes, each node has a color denoted by the array color[](color[i] denotes the color of ith node in form of an integer). Respond to ‘Q’ queries of the following type:

**distinct**u – Print the number of distinct colored nodes under the subtree rooted under ‘u’

Examples:

1 / \ 2 3 /|\ | \ 4 5 6 7 8 /| \ 9 10 11 color[] = {0, 2, 3, 3, 4, 1, 3, 4, 3, 2, 1, 1} Indexes NA 1 2 3 4 5 6 7 8 9 10 11 (Node Values and colors start from index 1) distinct 3 -> output should be '4'. There are six different nodes in subtree rooted with 3, nodes are 3, 7, 8, 9, 10 and 11. These nodes have four distinct colors (3, 4, 2 and 1) distinct 2 -> output should be '3'. distinct 7 -> output should be '3'.

Building a solution in steps:

- Flatten the tree using DFS; store the visiting time and ending time for every node in two arrays,
**vis_time[i]**stores the visiting time of the ith node while**end_time[i]**stores the ending time. - In the same DFS call, store the value of color of every node in an array
**flat_tree[]**, at indices:**vis_time[i]**and**end_time[i]**for ith node.

Note: size of the array**flat_tree[]**will be 2n.

Now the problem is reduced to finding the number of distinct elements in the range **[vis_time[u], end_time[u] ]** in the array **flat_tree[]** for each query of the specified type. To do so, we will process the queries off-line(processing the queries in an order different than the one provided in the question, and storing the results, and finally printing the result for each in the order specified in the question).

**Steps:**

- First, we pre-process the array
**flat_tree[]**; we maintain a**table[]**(an array of vectors),**table[i]**stores the vector containing all the indices in**flat_tree[]**that have value i. That is, if**flat_tree[j] = i**, then**table[i]**will have one of its element**j**. - In BIT, we update ‘1’ at ith index if we want the ith element of
**flat_tree[]**to be counted in**query()**method. We now maintain another array**traverser[]**;**traverser[i]**contains the pointer to the next element of table[i] that is not marked in BIT yet. - We now update our BIT and set ‘1’ at first occurrence of every element in
**flat_tree[]**and increment corresponding**traverser[]**by ‘1’(if**flat_tree[i]**is occurring for the first time then**traverser[flat_tree[i]]**is incremented by ‘1’) to point to the next occurrence of that element. - Now our
**query(R)**function for BIT would return the number of distinct elements in**flat_tree[]**in the range**[1, R]**. - We sort all the queries in order of increasing
**vis_time[]**, let**l**denote_{i}**vis_time[i]**and**r**denote the_{i}**end_time[i]**. Sorting the queries in increasing order of**l**gives us an edge, as when processing the ith query we won’t see any query in future with its ‘_{i}**l**‘ smaller than**l**. So we can remove all the elements’ occurrences up to_{i}**l**from BIT and add their next occurrences using the_{i}– 1**traverser[]**array. And then**query(R)**would return the number of distinct elements in the range**[l**._{i}, r_{i}]

// A C++ program implementing the above design #include<bits/stdc++.h> #define max_color 1000005 #define maxn 100005 using namespace std; // Note: All elements of global arrays are // initially zero // All the arrays have been described above int bit[maxn], vis_time[maxn], end_time[maxn]; int flat_tree[2 * maxn]; vector<int> tree[maxn]; vector<int> table[max_color]; int traverser[max_color]; bool vis[maxn]; int tim = 0; //li, ri and index are stored in queries vector //in that order, as the sort function will use //the value li for comparison vector< pair< pair<int, int>, int> > queries; //ans[i] stores answer to ith query int ans[maxn]; //update function to add val to idx in BIT void update(int idx, int val) { while ( idx < maxn ) { bit[idx] += val; idx += idx & -idx; } } //query function to find sum(1, idx) in BIT int query(int idx) { int res = 0; while ( idx > 0 ) { res += bit[idx]; idx -= idx & -idx; } return res; } void dfs(int v, int color[]) { //mark the node visited vis[v] = 1; //set visiting time of the node v vis_time[v] = ++tim; //use the color of node v to fill flat_tree[] flat_tree[tim] = color[v]; vector<int>::iterator it; for (it=tree[v].begin(); it!=tree[v].end(); it++) if (!vis[*it]) dfs(*it, color); // set ending time for node v end_time[v] = ++tim; // setting its color in flat_tree[] again flat_tree[tim] = color[v]; } //function to add an edge(u, v) to the tree void addEdge(int u, int v) { tree[u].push_back(v); tree[v].push_back(u); } //function to build the table[] and also add //first occurrences of elements to the BIT void hashMarkFirstOccurences(int n) { for (int i = 1 ; i <= 2 * n ; i++) { table[flat_tree[i]].push_back(i); //if it is the first occurence of the element //then add it to the BIT and increment traverser if (table[flat_tree[i]].size() == 1) { //add the occurence to bit update(i, 1); //make traverser point to next occurence traverser[flat_tree[i]]++; } } } //function to process all the queries and store thier answers void processQueries() { int j = 1; for (int i=0; i<queries.size(); i++) { //for each query remove all the ocurences before its li //li is the visiting time of the node //which is stored in first element of first pair for ( ; j < queries[i].first.first ; j++ ) { int elem = flat_tree[j]; //update(i, -1) removes an element at ith index //in the BIT update( table[elem][traverser[elem] - 1], -1); //if there is another occurrence of the same element if ( traverser[elem] < table[elem].size() ) { //add the occurrence to the BIT and //increment traverser update(table[elem][ traverser[elem] ], 1); traverser[elem]++; } } //store the answer for the query, the index of the query //is the second element of the pair //And ri is stored in second element of the first pair ans[queries[i].second] = query(queries[i].first.second); } } // Count distinct colors in subtrees rooted with qVer[0], // qVer[1], ...qVer[qn-1] void countDistinctColors(int color[], int n, int qVer[], int qn) { // build the flat_tree[], vis_time[] and end_time[] dfs(1, color); // add query for u = 3, 2 and 7 for (int i=0; i<qn; i++) queries.push_back(make_pair(make_pair(vis_time[qVer[i]], end_time[qVer[i]]), i) ); // sort the queries in order of increasing vis_time sort(queries.begin(), queries.end()); // make table[] and set '1' at first occurences of elements hashMarkFirstOccurences(n); // process queries processQueries(); // print all the answers, in order asked // in the question for (int i=0; i<queries.size() ; i++) { cout << "Distinct colors in the corresponding subtree" "is: " << ans[i] << endl; } } //driver code int main() { /* 1 / \ 2 3 /|\ | \ 4 5 6 7 8 /| \ 9 10 11 */ int n = 11; int color[] = {0, 2, 3, 3, 4, 1, 3, 4, 3, 2, 1, 1}; // add all the edges to the tree addEdge(1, 2); addEdge(1, 3); addEdge(2, 4); addEdge(2, 5); addEdge(2, 6); addEdge(3, 7); addEdge(3, 8); addEdge(7, 9); addEdge(7, 10); addEdge(7, 11); int qVer[] = {3, 2, 7}; int qn = sizeof(qVer)/sizeof(qVer[0]); countDistinctColors(color, n, qVer, qn); return 0; }

Output:

Distinct colors in the corresponding subtree is:4 Distinct colors in the corresponding subtree is:3 Distinct colors in the corresponding subtree is:3

Time Complexity:O( Q * log(n) )

This article is contributed by **Saumye Malhotra**. If you like GeeksforGeeks and would like to contribute, you can also write an article using contribute.geeksforgeeks.org or mail your article to contribute@geeksforgeeks.org. See your article appearing on the GeeksforGeeks main page and help other Geeks.

Please write comments if you find anything incorrect, or you want to share more information about the topic discussed above.

**Practice Tags :**