620E New Year Tree

题意

给定一棵树, 根节点为1, 每个节点上有颜色\(c_i\) , \(m\)次操作
操作1: 1 u c, 把u为根的子树上的所有节点修改为c
操作2: 2 u, 查询以u为根的子树的所有节点的颜色种类的数量

思路

新知识点:dfs序
通过in, out, pos, 可以将树转换为区间

代码

#include<bits/stdc++.h>
using namespace std;
#define int long long
const int N = 4e5+10;

struct Node{
    int l, r, c;
    int cover;
};
vector<int> ed[N];
int c[N];

Node tr[N << 2];
int tot = 0;
int in[N], out[N], pos[N];

void dfs(int u, int fa)
{
    pos[++tot] = u;
    // cout << tot << " " << u << endl;
    in[u] = tot;
    for(auto e : ed[u])
    {
        if(e == fa)continue;
        dfs(e, u);
    }
    out[u] = tot;
}
inline int ls(int p){return p << 1;}
inline int rs(int p){return p << 1 | 1;}

inline void push_up(int p)
{
    tr[p].c = tr[ls(p)].c | tr[rs(p)].c;
}

void build(int p, int l, int r)
{
    tr[p] = {l, r};
    if(l == r)
    {
        tr[p].c = tr[p].cover = 1ll << c[pos[l]];
        return ;
    }
    int m = l + r >> 1;
    build(ls(p), l, m);
    build(rs(p), m + 1, r);
    push_up(p);
}
void f(int cd, int fa)
{
    tr[cd].cover = tr[fa].cover;
    tr[cd].c = tr[fa].cover;
}
void push_down(int p)
{
    if(tr[p].cover)
    {
        f(ls(p), p);
        f(rs(p), p);
        tr[p].cover = 0;
    }
}
void change(int p, int l, int r, int c)
{
    if(tr[p].l >= l && tr[p].r <= r)
    {
        tr[p].cover = 1ll<<c;
        tr[p].c = 1ll<<c;
        return ;
    }
    push_down(p);
    int m = tr[p].l + tr[p].r >> 1;
    if(l <= m) change(ls(p), l, r, c);
    if(r > m) change(rs(p), l, r, c);
    push_up(p);
}
int Q(int p, int l, int r)
{
    if(tr[p].l >= l && tr[p].r <= r)
    {
        return tr[p].c;
    }
    push_down(p);
    int c = 0;
    int m = tr[p].l + tr[p].r >> 1;
    if(l <= m)
        c |= Q(ls(p), l, r);
    if(r > m)
        c |= Q(rs(p), l, r);
    return c;
}
signed main()
{
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int n, m;
    cin >> n >> m;
    
    for(int i = 1; i <= n; i++)
    {
        cin >> c[i];
    }
    
    for(int i = 1; i <= n - 1; i++)
    {
        int x, y;
        cin >> x >> y;
        ed[x].push_back(y);
        ed[y].push_back(x);
    }
    dfs(1, -1);
    
    build(1, 1, n);
    for(int i = 1; i <= m; i++)
    {
        int op;
        cin >> op;
        if(op == 1)
        {
            int u, c;
            cin >> u >> c;
            change(1, in[u], out[u], c);
            
        }
        else{
            int u;
            cin >> u;
            auto c = Q(1, in[u], out[u]);
            int ans = __builtin_popcountll(c);
            cout << ans << '\n';
        }
    }
    return 0;
}