Shared code

#include<bits/stdc++.h>
using namespace std;
 
#define sum(a,b,m) ((a%m)+(b%m))%m
#define diff(a,b,m) ((a%m)-(b%m)+m)%m
#define pro(a,b,m) ((a%m)*(b%m))%m
 
 
#define ll long long 
#define nl "\n"
#define M 1000000007

class SegmentTree{
    int n;
    set<int> s;
    vector<int> arr;
    vector<set<int>> tree;

    public:

        SegmentTree(int n,vector<int> &arr){
            s.clear();
            this->n=n;
            this->arr=arr;
            tree.resize(4*n);
        }

        set<int> merge(set<int> a,set<int> b){
            set<int> c=a;
            for(auto &p:b){
                c.insert(p);
            }
            return c;
        }
        void build(int start,int end,int ind){
            if(start==end){
                tree[ind].insert(arr[start]);
                return;
            }
            int mid=(start+end)>>1;
            build(start,mid,2*ind);
            build(mid+1,end,2*ind+1);
            tree[ind]=merge(tree[2*ind],tree[2*ind+1]);
        }

        set<int> query(int l,int r,int start,int end,int ind){
            if(l<=start && r>=end){
                return tree[ind];
            }
            else if(l>end || r<start){
                return this->s;
            }
            int mid=(start+end)>>1;
            return merge(query(l,r,start,mid,2*ind),query(l,r,mid+1,end,2*ind+1));
        }
};
class Graph{
    int n,timer;
    vector<vector<int>> adj;
    vector<int> color;
    vector<int> intime,outtime;
    vector<int> flat;

    public:
        Graph(int n,vector<int> color){
            this->n=n;
            this->color.resize(n);
            this->color=color;
            adj.clear();
            flat.clear();
            adj.resize(n+1);
            intime.resize(n+1);
            outtime.resize(n+1);
            timer=0;
        }
        
        void add_edge(int u,int v){
            adj[u].push_back(v);
            adj[v].push_back(u);
        }

        void dfs(int src,int par){
            intime[src]=timer;
            timer++;
            flat.push_back(color[src-1]);
            for(auto &child:adj[src]){
                if(child!=par){
                    dfs(child,src);
                }
            }
            outtime[src]=timer;
            timer++;
            flat.push_back(color[src-1]);
        }

        void solve(){
            dfs(1,-1);
            SegmentTree sg=SegmentTree(2*n,flat);
            sg.build(0,2*n-1,1);
            for(int i=1;i<=n;i++){
                cout<<sg.query(intime[i],outtime[i],0,2*n-1,1).size()<<" ";
            }
            cout<<nl;
        }
};
void solve(){
    int n;cin>>n;
    vector<int> color(n);
    for(auto &p:color){
        cin>>p;
    }
    Graph g=Graph(n,color);
    for(int i=1;i<n;i++){
        int u,v;cin>>u>>v;
        g.add_edge(u,v);
    }
    g.solve();
}

int main(){
    int t=1;
    while(t--){
        solve();
    }
}