Z - New Year Tree CodeForces - 620E 线段树 区间种类 bitset

#include <cstring>
#include <queue>
#include <cstdlib>
#include <cstdio>
#include <iostream>
#include <string>
#include <bitset>
#include <algorithm>
#include <map>
#include <vector>
#define inf 0x3f3f3f3f
#define inf64 0x3f3f3f3f3f3f3f3f
using namespace std;
typedef long long ll;
typedef bitset<100> bit;
const int maxn = 4e5 + 10;
vector<int>G[maxn];
ll el[maxn], er[maxn], a[maxn], tot, num[maxn];

void dfs(int u,int pre)
{
    el[u] = ++tot;
    num[tot] = u;
    for(int i=0;i<G[u].size();i++)
    {
        int v = G[u][i];
        if (v == pre) continue;
        dfs(v, u);
    }
    er[u] = tot;
}

bit tree[maxn * 8];
ll lazy[maxn * 8];

void build(int id,int l,int r)
{
    lazy[id] = 0;
    if(l==r)
    {
        tree[id].reset();
        tree[id] = (1ll << (a[num[l]] - 1));
        return;
    }
    int mid=(l + r) >> 1;
    build(id << 1, l, mid);
    build(id << 1 | 1, mid + 1, r);
    tree[id] = tree[id << 1] | tree[id << 1 | 1];
}

void push_down(int id)
{
    if(lazy[id]!=0)
    {
        tree[id << 1].reset(); tree[id << 1 | 1].reset();
        tree[id << 1] = (1ll << (lazy[id] - 1));
        tree[id << 1 | 1] = (1ll << (lazy[id] - 1));
        lazy[id << 1] = lazy[id << 1 | 1] = lazy[id];
        lazy[id] = 0;
    }
}

void update(int id,int l,int r,int x,int y,int val)
{
    if(x<=l&&y>=r)
    {
        tree[id].reset();
        tree[id] = (1ll << (val - 1));
        lazy[id] = val;
        return;
    }
    push_down(id);
    int mid = (l + r) >> 1;
    if (x <= mid) update(id << 1, l, mid, x, y, val);
    if (y > mid) update(id << 1 | 1, mid + 1, r, x, y, val);
    tree[id] = tree[id << 1] | tree[id << 1 | 1];
}

bit query(int id,int l,int r,int x,int y)
{
    if (x <= l && y >= r) return tree[id];
    int mid = (l + r) >> 1;
    bit ans; 
    ans.reset();
    push_down(id);
    if (x <= mid) ans = query(id << 1, l, mid, x, y);
    if (y > mid) ans |= query(id << 1 | 1, mid + 1, r, x, y);
    return ans;
}

int main() {
    int n, m;
    tot = 0;
    scanf("%d%d", &n, &m);
    for (int i = 1; i <= n; i++) scanf("%lld", &a[i]);
    for (int i = 1; i < n; i++) {
        int u, v;
        scanf("%d%d", &u, &v);
        G[u].push_back(v);
        G[v].push_back(u);
    }
    dfs(1, -1);
    build(1, 1, n);
    while (m--) {
        int opt, v, x;
        scanf("%d", &opt);
        if (opt == 1) {
            scanf("%d%d", &v, &x);
            update(1, 1, n, el[v], er[v], x);
        }
        else {
            scanf("%d", &v);
            bit ans = query(1, 1, n, el[v], er[v]);
            printf("%d\n", ans.count());
        }
    }
    return 0;
}