Problem Statement : 620E - New Year Tree My Solution :

#include<bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
#include <functional>
#define PIKACHU {ios_base::sync_with_stdio(0); cin.tie(0);} 
#define flush(a) for(auto ptr: a) cout<<ptr<<" "; cout<<endl;
#define flushh(a,s,e) for(int i=s;i<=e;i++) cout<<a[i]<<" "; cout<<endl;
#define int long long int 
#define all(a) a.begin(),a.end() 
#define pb push_back
#define vvi vector<vector<int>>
#define vi vector<int>
#define pii pair<int,int>
#define read(a,n) for(int i=0;i<n;i++) cin>>a[i];
#define ff first
#define ss second
#define yes {cout<<"YES\n"; return;}
#define no {cout<<"NO\n"; return;}
#define mod 1000000007
#define modd 998244353
#define endl '\n'
#define watch(x) cout<<"The value of "<<(#x)<<" is "<<(x)<<endl
#define f_tree tree<int, null_type, less<int>, rb_tree_tag,  tree_order_statistics_node_update>
using namespace __gnu_pbds;
using namespace std;

//binary search maybe...

const int N=1200300;
vector<int> adj[N];
vector<int> euler;
vector<int> bitmask(N);
vi color(N);
map<int,pii> mp;
int t[N];
int lim;

void dfs(int node, int par=0)
{
euler.pb(node);
for(auto child: adj[node])
{
    if(child!=par)
    dfs(child,node);
}
euler.pb(node);
}

void build(int id, int l, int r)
{
    if(l==r) {t[id]=color[l]; return;}
    int mid=(l+r)/2;
    build(2*id,l,mid);
    build(2*id+1,mid+1,r);
    t[id]=(t[2*id]|t[2*id+1]);
}

void update(int id, int l, int r, int pos, int value)
{
    if(pos<l || pos>r) return;
    if(l==r) {t[id]=value; return;}
    int mid=(l+r)/2;
    update(2*id,l,mid,pos,value);
    update(2*id+1,mid+1,r,pos,value);
    t[id]=(t[2*id]|t[2*id+1]);
}

int query(int id, int l, int r, int lq, int rq)
{
    if(l>rq || r<lq) return 0;
    if(lq<=l && r<=rq) return t[id];
    int mid=(l+r)/2;
    return (query(2*id,l,mid,lq,rq)|query(2*id+1,mid+1,r,lq,rq));
}

void solve()
{
    int n,q; cin>>n>>q;
    for(int i=1;i<=n;i++) {int x; cin>>x; bitmask[i]=(1ll<<x);}
    int m=n-1; while(m--)
    {
        int u,v; cin>>u>>v;
        adj[u].pb(v),adj[v].pb(u);
    } 
    dfs(1);
    lim=euler.size();
    vi mm(n+1,0);
    for(int i=0;i<euler.size();i++)
    {
        mm[euler[i]]++;
        if(mm[euler[i]]==1) mp[euler[i]].ff=i,color[i]=bitmask[euler[i]];
        else mp[euler[i]].ss=i,color[i]=0;
    }
    build(1,0,lim-1);
    while(q--)
    {
        int type; cin>>type;
        if(type==1)
        {
            int s,x; cin>>s>>x;
            update(1,0,lim-1,mp[s].ff,(1ll<<x));
            update(1,0,lim-1,mp[s].ss,(1ll<<x));
        }
        else
        {
            int s;cin>>s;
            int ans=query(1,0,lim-1,mp[s].ff,mp[s].ss);
            cout<<__builtin_popcount(ans)<<endl;
        }
    }
}

int32_t main()
{
    PIKACHU;
    //int _; cin>>_; while(_--)
    solve();
    return 0;
}

Can anyone spot the mistake ? I'd be thankful.