Coloring Tree
Difficulty: Hard



You are given a tree with **N** nodes with every node being colored. A color is represented by an integer ranging from 1 to 10<sup>9</sup>. Can you find the number of distinct colors available in a subtree rooted at the node **s**? 

**Input Format**  
The first line contains three space separated integers representing the number of nodes in the tree (**N**), number of queries to answer (**M**) and the root of the tree. 

In each of the next N-1 lines, there are two space separated integers(a b) representing an edge from node a to Node b and vice-versa.   

N lines follow: N+i<sup>th</sup> line contains the color of the i<sup>th</sup> node.

M lines follow: Each line containg a single integer s.

**Output Format**  
Output exactly M lines, each line containing the output of the i<sub>th</sub> query.

**Constraints**  
0 <= M <= 10<sup>5</sup><br>
1 <= N <= 10<sup>5</sup><br>
1 <= root <= N<br>
1 <= color of the Node <= 10<sup>9</sup><br>

**Example**

**Sample Input**

	4 2 1
    1 2
    2 4
    2 3
    10
    20
    20
    30
    1
    2

**Sample Output**

	3
    2

**Explanation**

Query 1-Subtree rooted at 1 is the entire tree and colors used are 10 20 20 30 , so the answer is 3(10,20 and 30)

Query 2-Subtree rooted at 2 contains color 20 20 30, so the answer is 2(20 and 30)


Input Format:

Output Format:

Sample Input:

Sample Output:

Python 3 starter code:
#
# Complete the 'solve' function below.
#
# The function is expected to return an INTEGER_ARRAY.
# The function accepts following parameters:
#  1. 2D_INTEGER_ARRAY tree
#  2. INTEGER_ARRAY color
#  3. INTEGER_ARRAY s
#

def solve(tree, color, s):
    # Write your code here



Editorial: Coloring Tree - 101 Hack February Editorial
<style id="MathJax_SVG_styles">.MathJax_SVG_Display {text-align: center; margin: 1em 0em; position: relative; display: block!important; text-indent: 0; max-width: none; max-height: none; min-width: 0; min-height: 0; width: 100%}
.MathJax_SVG .MJX-monospace {font-family: monospace}
.MathJax_SVG .MJX-sans-serif {font-family: sans-serif}
.MathJax_SVG {display: inline; font-style: normal; font-weight: normal; line-height: normal; font-size: 100%; font-size-adjust: none; text-indent: 0; text-align: left; text-transform: none; letter-spacing: normal; word-spacing: normal; word-wrap: normal; white-space: nowrap; float: none; direction: ltr; max-width: none; max-height: none; min-width: 0; min-height: 0; border: 0; padding: 0; margin: 0}
.MathJax_SVG * {transition: none; -webkit-transition: none; -moz-transition: none; -ms-transition: none; -o-transition: none}
.mjx-svg-href {fill: blue; stroke: blue}
</style><svg style="display: none;"><defs id="MathJax_SVG_glyphs"></defs></svg><p><strong>Difficulty Level</strong> : Medium <br>
<strong>Required Knowledge</strong> : DFS and Segment Trees<br>
<strong>Problem Setter</strong> : <a href="https://www.hackerrank.com/devuy11">Devendra Agarwal</a>  <br>
<strong>Problem Tester</strong> : <a href="https://www.hackerrank.com/darkshadows">Lalit Kundu</a>  <br></p>

<p>Store the order in which you execute DFS starting from the given root.Now observe that each subtree will correspond to the sub-array of the array formed by DFS. Store the starting and ending position for each subtree in the array.</p>

<p>Now the question reduces to finding the number of distinct value in the sub-array.</p>

<p>Note that atmax 10^5 different colors are used as N&lt;=10^5 ,so store the colors used in the tree and map it to an integer from 1 to 10^5.Now the question is to find number of distinct values in the sub-array given that all the values are from 1 to 10^5.</p>

<p>The result of a query [a, b] is number of integers whose last occurrence in [1, b] is &gt;= a(Last Occurence in [1,b] is in [a,b]).</p>

<p>Let's have two kinds of events: QueryEvent and NumberEvent. First we read whole input and sort all events, the key for QueryEvents are their end (i.e. b for query [a, b]), and for NumberEvents the key is their position in array.</p>

<p>We also have a segment tree which answers the queries of kind: how many elements are at position range [x, y]. Initially the tree is empty.</p>

<p>Then we process events in order:<br>
<strong>a</strong>. When we meet a NumberEvent:<br>
  1. If number has occurred before ,then we put 0 in that position of seg tree and update the tree.<br>
    2. We put 1 in position of number in the segment tree.<br>
<strong>b</strong>.When we meet a QueryEvent:<br>
    - Query the segment tree (for finding number of 1's in the range), and find the answer.<br></p>

<p><strong>Setter's Code</strong>  </p>

<pre><code>#include&lt;stdio.h&gt;
#include&lt;stdlib.h&gt;
#include&lt;iostream&gt;
#include&lt;algorithm&gt;
#include&lt;vector&gt;
#include&lt;string.h&gt;
using namespace std;
#define Max_N 1000001
#define INF 1000000000
int N,color[Max_N],root,from,to,counter,parent[Max_N],s,q,search_here[Max_N],highest,last_occurence[Max_N],where,val,ind ,seg[1&lt;&lt;21],power[25],level,answer[Max_N];
vector&lt;int&gt; myvector[Max_N];
struct transfer
{
    int start,end;
}Node[Max_N];
typedef struct queries qu;
struct queries
{
    int ind ,left,right;
    bool operator&lt;(const qu q)const{
        return right&lt;q.right;
    }
}Q[Max_N];
void DFS(int rooti)
{
    Node[rooti].start=counter;
    counter++;
    while(!myvector[rooti].empty()){
        to=myvector[rooti].back();
        if(to!=parent[rooti]){
            parent[to]=rooti;
            DFS(to);
        }
        myvector[rooti].pop_back();
    }
    Node[rooti].end=counter-1;
}
int search(int low,int high,int s)
{
    int mid=(low+high)&gt;&gt;1;
    if(search_here[mid]==s) return mid;
    else if(search_here[mid]&gt;s) return search(low,mid-1,s);
    else    return search(mid+1,high,s);
}
void constant()
{
    power[0]=1;
    for(int i=1;i&lt;=23;i++)  power[i]=power[i-1]&lt;&lt;1;
}
void obtain(int x)
{
    int y=x;
    level=0;
    while(x){
        x&gt;&gt;=1;
        level++;
    }
    if(power[level-1]!=y)   level++;
}
void update(int pos,int val)
{
    seg[pos]=val;
    pos=pos&gt;&gt;1;
    while(pos){
        seg[pos]=seg[pos&lt;&lt;1]+seg[(pos&lt;&lt;1)+1];
        pos&gt;&gt;=1;
    }
    return;
}
int Query(int s,int e,int i,int j,int node)
{
    if(i&gt;e || s&gt;j)  return 0;
    if(s&gt;=i &amp;&amp; j&gt;=e)    return seg[node];
    return (Query(s,((s+e)&gt;&gt;1),i,j,(node&lt;&lt;1))+Query(((s+e)&gt;&gt;1)+1,e,i,j,(node&lt;&lt;1)+1));
}
int main()
{
    int v;
    scanf("%d%d%d",&amp;N,&amp;q,&amp;root);
    //connections
    for(int i=1;i&lt;N;i++){
        scanf("%d%d",&amp;from,&amp;to);
        myvector[from].push_back(to);
        myvector[to].push_back(from);
    }
    counter=1;
    parent[root]=root;
    DFS(root);
    //color of the graph
    for(int i=1;i&lt;=N;i++){
        scanf("%d",&amp;v);
        color[Node[i].start]=v;
        search_here[i]=v;   //finally i need to sort so no troubles
    }


    for(int i=0;i&lt;q;i++){
        scanf("%d",&amp;s);
        Q[i].ind =i;
        Q[i].left=Node[s].start;
        Q[i].right=Node[s].end;
    }
    sort(Q,Q+q);
    sort(search_here+1,search_here+N+1);
    highest=2;
    for(int i=2;i&lt;=N;i++){  
        if(search_here[i]==search_here[highest-1])  continue;
        else{
            search_here[highest]=search_here[i];
            highest++;
        }
    }
    highest--;
    where=0;
    constant();
    obtain(N);
    for(int i=0;i&lt;=N;i++)   last_occurence[i]=-INF;
    memset(seg,0,sizeof(seg));
    //search in search_here from 1 to highest ind ed
    for(int i=1;i&lt;=N;i++){
        val=color[i];
        ind =search(1,highest,val);
        if(last_occurence[ind ]&lt;0){
            last_occurence[ind]=i;
            update(power[level-1]+last_occurence[ind]-1,1);
        }
        else{
            update(power[level-1]-1+last_occurence[ind ],0);
            last_occurence[ind]=i;
            update(power[level-1]-1+last_occurence[ind ],1);
        }
        while(Q[where].right==i){
            answer[Q[where].ind]=Query(1,power[level-1],Q[where].left,Q[where].right,1);
            where++;
        }
    }
    for(int i=0;i&lt;q;i++)    printf("%d\n",answer[i]);
    return 0;
}
</code></pre>

<p><strong>Tester's Code</strong>  </p>

<pre><code>/*
Author:
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%   LALIT KUNDU      %%%%%%%% 
%%%%%%%%   IIIT HYDERABAD   %%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
*/
#include&lt;stdio.h&gt;
#include&lt;stdlib.h&gt;
#include&lt;string.h&gt;
#include&lt;math.h&gt;
#include&lt;iostream&gt;
#include&lt;vector&gt;
#include&lt;cassert&gt;
#include&lt;sstream&gt;
#include&lt;map&gt;
#include&lt;set&gt;
#include&lt;stack&gt;
#include&lt;queue&gt;
#include&lt;algorithm&gt;
using namespace std;
#define pb push_back
#define mp make_pair
#define clr(x) x.clear()
#define sz(x) ((int)(x).size())
#define F first
#define S second
#define REP(i,a,b) for(i=a;i&lt;b;i++)
#define rep(i,b) for(i=0;i&lt;b;i++)
#define rep1(i,b) for(i=1;i&lt;=b;i++)
#define mod 1000000007
#define pdn(n) printf("%d\n",n)
#define sl(n) scanf("%d",&amp;n)
#define sd(n) scanf("%d",&amp;n)
#define pn printf("\n")
typedef pair&lt;int,int&gt; PII;
typedef vector&lt;PII&gt; VPII;
typedef vector&lt;int&gt; VI;
typedef vector&lt;VI&gt; VVI;
typedef long long LL;
#define INF 1000000007
int LMITR[1000009];
map &lt; int ,int &gt; PI;
#define MAX 10000009
int tree[MAX+1]={};
void update(int idx, int val) {
    while (idx &lt;= MAX) {
        tree[idx] += val;
        idx += (idx &amp; -idx);
    }
}
int readi(int idx) {
    int sum = 0;
    while (idx &gt; 0) {
        sum +=(int) tree[idx];
        idx -= (idx &amp; -idx);
    }
    return sum;
}
int to[1000005],ar[1000009];
void foo(int n, int q, vector &lt; pair &lt; pair &lt; int,int &gt;, int &gt;  &gt;  qu)
{
    int i;
    for(i=n; i&gt;=1; i--)
    {
        if(PI.count(ar[i])==0)
            LMITR[i]=INF;
        else LMITR[i]=PI[ar[i]];
        PI[ar[i]]=i;
    }
    PI.clear();
    for(i=1; i&lt;=n; i++)
        if(PI.count(ar[i])==0)
            PI[ar[i]]=1,update(i,1);
    sort(qu.begin(),qu.end());
    int cur=0,ans[1000007];
    for(i=1; i&lt;=n; i++)
    {
        while(cur&lt;q &amp;&amp; qu[cur].F.F==i)
        {
            ans[qu[cur].S]=readi(qu[cur].F.S);
            cur++;
        }
        if(LMITR[i]!=INF)
            update(LMITR[i],1);
        if(i==1)
        {
            if(readi(1))
                update(1,-readi(1));
        }
        else
            update(i,readi(i-1)-readi(i));
    }
    for(i=0; i&lt;q; i++)
        printf("%d\n",ans[i]);
}
VI arr[1000009];
int mymap[1000009]={},counter=0,flag[1000009]={},siz[1000009]={};
string str;
int dfs1(int p)
{
    flag[p]=1;
    if(siz[p]!=-1)return siz[p];
    int i,ret=1;
    for(i=0; i&lt;arr[p].size(); i++)
        if(flag[arr[p][i]]==0)ret+=dfs1(arr[p][i]);
    siz[p]=ret;
    return ret;
}
void dfs(int p)
{
    int j,i;
    stack &lt; int &gt; mystack;
    mystack.push(p);
    flag[p]=1;
    while(!mystack.empty())
    {
        i=mystack.top();
        mystack.pop();
        mymap[i]=counter;
        counter++;
        for(int j=0;j&lt;arr[i].size();j++)
        {
            if(flag[arr[i][j]]==0)
                mystack.push(arr[i][j]),flag[arr[i][j]]=1;
        }
    }
}
int temp[1000009]={};
int main()
{
    int n,q,i,j,u,v,root,k;
    sl(n),sl(q),sl(root);
    root--;
    for(i=1; i&lt;n; i++)
    {
        sl(u),sl(v);
        u--,v--;
        arr[u].pb(v);
        arr[v].pb(u);
    }
    for(i=0; i&lt;n; i++)
        sl(temp[i]);
    dfs(root);
    memset(siz,-1,sizeof(siz));
    memset(flag,0,sizeof(flag));
    dfs1(root);
    for(i=1; i&lt;=n; i++)
    {
        ar[mymap[i-1]+1]=temp[i-1];
    }
    vector &lt; pair &lt; pair &lt; int ,int &gt; , int &gt; &gt; query;
    for(i=0; i&lt;q; i++)
    {
        scanf("%d",&amp;j);
    query.pb(mp(mp(mymap[j-1]+1,mymap[j-1]+siz[j-1]),i));

    }
    foo(n,q,query);
    return 0;
}
</code></pre>

Reference solution (C++):

    #include<stdio.h>
    #include<stdlib.h>
    #include<iostream>
    #include<algorithm>
    #include<vector>
    #include<string.h>
    using namespace std;
    #define Max_N 1000001
    #define INF 1000000000
    int N,color[Max_N],root,from,to,counter,parent[Max_N],s,q,search_here[Max_N],highest,last_occurence[Max_N],where,val,ind ,seg[1<<21],power[25],level,answer[Max_N];
    vector<int> myvector[Max_N];
    struct transfer
    {
    	int start,end;
    }Node[Max_N];
    typedef struct queries qu;
    struct queries
    {
    	int ind ,left,right;
    	bool operator<(const qu q)const{
    		return right<q.right;
    	}
    }Q[Max_N];
    void DFS(int rooti)
    {
    	Node[rooti].start=counter;
    	counter++;
    	while(!myvector[rooti].empty()){
    		to=myvector[rooti].back();
    		if(to!=parent[rooti]){
    			parent[to]=rooti;
    			DFS(to);
    		}
    		myvector[rooti].pop_back();
    	}
    	Node[rooti].end=counter-1;
    }
    int search(int low,int high,int s)
    {
    	int mid=(low+high)>>1;
    	if(search_here[mid]==s)	return mid;
    	else if(search_here[mid]>s)	return search(low,mid-1,s);
    	else	return search(mid+1,high,s);
    }
    void constant()
    {
    	power[0]=1;
    	for(int i=1;i<=23;i++)	power[i]=power[i-1]<<1;
    }
    void obtain(int x)
    {
    	int y=x;
    	level=0;
    	while(x){
    		x>>=1;
    		level++;
    	}
    	if(power[level-1]!=y)	level++;
    }
    void update(int pos,int val)
    {
    	seg[pos]=val;
    	pos=pos>>1;
    	while(pos){
    		seg[pos]=seg[pos<<1]+seg[(pos<<1)+1];
    		pos>>=1;
    	}
    	return;
    }
    int Query(int s,int e,int i,int j,int node)
    {
    	if(i>e || s>j)	return 0;
    	if(s>=i && j>=e)	return seg[node];
    	return (Query(s,((s+e)>>1),i,j,(node<<1))+Query(((s+e)>>1)+1,e,i,j,(node<<1)+1));
    }
    int main()
    {
    	int v;
    	scanf("%d%d%d",&N,&q,&root);
    	//connections
    	for(int i=1;i<N;i++){
    		scanf("%d%d",&from,&to);
    		myvector[from].push_back(to);
    		myvector[to].push_back(from);
    	}
    	counter=1;
    	parent[root]=root;
    	DFS(root);
    	//color of the graph
    	for(int i=1;i<=N;i++){
    		scanf("%d",&v);
    		color[Node[i].start]=v;
    		search_here[i]=v;	//finally i need to sort so no troubles
    	}
    	
    	
    	for(int i=0;i<q;i++){
    		scanf("%d",&s);
    		Q[i].ind =i;
    		Q[i].left=Node[s].start;
    		Q[i].right=Node[s].end;
    	}
    	sort(Q,Q+q);
    	sort(search_here+1,search_here+N+1);
    	highest=2;
    	for(int i=2;i<=N;i++){	
    		if(search_here[i]==search_here[highest-1])	continue;
    		else{
    			search_here[highest]=search_here[i];
    			highest++;
    		}
    	}
    	highest--;
    	where=0;
    	constant();
    	obtain(N);
    	for(int i=0;i<=N;i++)	last_occurence[i]=-INF;
    	memset(seg,0,sizeof(seg));
    	//search in search_here from 1 to highest ind ed
    	for(int i=1;i<=N;i++){
    		val=color[i];
    		ind =search(1,highest,val);
    		if(last_occurence[ind ]<0){
    			last_occurence[ind]=i;
    			update(power[level-1]+last_occurence[ind]-1,1);
    		}
    		else{
    			update(power[level-1]-1+last_occurence[ind ],0);
    			last_occurence[ind]=i;
    			update(power[level-1]-1+last_occurence[ind ],1);
    		}
    		while(Q[where].right==i){
    			answer[Q[where].ind]=Query(1,power[level-1],Q[where].left,Q[where].right,1);
    			where++;
    		}
    	}
    	for(int i=0;i<q;i++)	printf("%d\n",answer[i]);
    	return 0;
    }