[Codeforces] E. New Year Tree

Author            : Dipu Kumar Mohanto 
                    CSE, Batch - 6
                    BRUR.
Problem Statement : E. New Year Tree
Source            : Codeorces
Category          : Graph Theory, Data Structure
Algorithm         : Euler tour on Tree, Segment Tree, Lazy Propagation
Verdict           : Accepted

#include "bits/stdc++.h"
 
using namespace std;
 
#define ll                     long long
#define eb                     emplace_back
 
static const int maxn = 4e5 + 5;
 
struct info
{
    ll mask, prop;
    info(ll mask = 0, ll prop = 0) : mask(mask), prop(prop) {}
} tree[maxn << 2];
 
vector <int> graph[maxn];
int tin[maxn], tout[maxn];
int dfstime, tnode, tquery, color[maxn];
 
void push(int node, int a, int b)
{
    if (tree[node].prop == 0 || a == b) return;
    int left = node << 1;
    int right = left | 1;
    tree[left] = tree[right] = info(1LL << tree[node].prop, tree[node].prop);
    tree[node].prop = 0;
}
 
void update(int node, int a, int b, int i, int j, int c)
{
    if (a > b || a > j || b < i) return;
    if (a >= i && b <= j)
    {
        tree[node] = info(1LL << c, c);
        return;
    }
    push(node, a, b);
    int left = node << 1;
    int right = left | 1;
    int mid = (a+b)>>1;
    update(left, a, mid, i, j, c);
    update(right, mid+1, b, i, j, c);
    tree[node].mask = tree[left].mask | tree[right].mask;
}
 
ll query(int node, int a, int b, int i, int j)
{
    if (a > b || a > j || b < i) return 0LL;
    push(node, a, b);
    if (a >= i && b <= j) return tree[node].mask;
    int left = node << 1;
    int right = left | 1;
    int mid = (a+b)>>1;
    ll p1 = query(left, a, mid, i, j);
    ll p2 = query(right, mid+1, b, i, j);
    return p1 | p2;
}
 
void dfs(int u = 1, int p = -1)
{
    dfstime++;
    tin[u] = dfstime;
    for (int v : graph[u])
    {
        if (v == p) continue;
        dfs(v, u);
    }
    tout[u] = dfstime;
    update(1, 1, tnode, tin[u], tin[u], color[u]);
}
 
int main()
{
    //freopen("in.txt", "r", stdin);
 
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    cout.tie(NULL);
 
    cin >> tnode >> tquery;
    for (int i = 1; i <= tnode; i++)
        cin >> color[i];
    for (int e = 1; e < tnode; e++)
    {
        int u, v;
        cin >> u >> v;
        graph[u].eb(v);
        graph[v].eb(u);
    }
    dfs();
    while (tquery--)
    {
        int type;
        cin >> type;
        if (type == 1)
        {
            int v, c;
            cin >> v >> c;
            update(1, 1, tnode, tin[v], tout[v], c);
        }
        else
        {
            int v;
            cin >> v;
            ll mask = query(1, 1, tnode, tin[v], tout[v]);
            int diff = __builtin_popcountll(mask);
            cout << diff << "\n";
        }
    }
}