CodeForces 620E New Year Tree【树链剖分】【状态压缩】【位运算】

板子题。

用一个 long long 记录一下 $60$ 种颜色是否有了，然后直接用树链剖分统计就可以了。

可能稍微有点卡常。

#include <bits/stdc++.h>
#define int long long
using namespace std;
const int N = 4e5 + 10;
int dep[N], fa[N];
int sz[N];
int cnt, w[N], dfs_order[N], rk[N], top[N], len[N];
int col[N];
vector<int> z[N];
int n, q;
void dfs1(int now, int fat = -1)
{
    dep[now] = dep[fa[now]] + 1;
    sz[now] = 1;
    for (auto &u : z[now])
        if (u != fat)
        {
            fa[u] = now;
            dfs1(u, now);
            sz[now] += sz[u];
        }
}
void dfs2(int now, int head, int fa = -1)
{
    cnt++;
    w[cnt] = 1ll << col[now];
    dfs_order[cnt] = now;
    rk[now] = cnt;
    top[now] = head;
    len[head]++;
    int p = -1;
    for (auto &u : z[now])
    {
        if (u == fa)
            continue;
        if (p == -1)
            p = u;
        else if (sz[u] > sz[p])
            p = u;
    }
    if (p == -1)
        return;
    dfs2(p, head, now);
    for (auto &j : z[now])
        if (p != j && j != fa)
            dfs2(j, j, now);
}

struct Node
{
    int l, r, sum, lazy;
    Node()
    {
        l = r = sum = lazy = 0;
    }
    ~Node()
    {
        l = r = sum = lazy = 0;
    }
    void init(int p)
    {
        l = r = p;
        sum = w[p];
        lazy = 0;
    }
    void color(int p)
    {
        sum = p;
        lazy = p;
    }
} y[N << 2];

Node operator+(const Node &lhs, const Node &rhs)
{
    Node res;
    res.l = lhs.l;
    res.r = rhs.r;
    res.sum = lhs.sum | rhs.sum;
    res.lazy = 0;
    return res;
}

void build(int l, int r, int rt)
{
    if (l == r)
        y[rt].init(l);
    else
    {
        int mid = l + r >> 1;
        build(l, mid, rt << 1);
        build(mid + 1, r, rt << 1 | 1);
        y[rt] = y[rt << 1] + y[rt << 1 | 1];
    }
}

void push_lazy(int rt)
{
    if (y[rt].lazy)
    {
        y[rt << 1].color(y[rt].lazy);
        y[rt << 1 | 1].color(y[rt].lazy);
        y[rt].lazy = 0;
    }
}

void modify(int l, int r, int rt, int nowl, int nowr, int value)
{
    if (nowl <= l && r <= nowr)
        y[rt].color(value);
    else
    {
        int mid = l + r >> 1;
        push_lazy(rt);
        if (nowl <= mid)
            modify(l, mid, rt << 1, nowl, nowr, value);
        if (mid < nowr)
            modify(mid + 1, r, rt << 1 | 1, nowl, nowr, value);
        y[rt] = y[rt << 1] + y[rt << 1 | 1];
    }
}

int query(int l, int r, int rt, int nowl, int nowr)
{
    if (nowl <= l && r <= nowr)
        return y[rt].sum;
    int mid = l + r >> 1;
    push_lazy(rt);
    int sum = 0;
    if (nowl <= mid)
        sum |= query(l, mid, rt << 1, nowl, nowr);
    if (mid < nowr)
        sum |= query(mid + 1, r, rt << 1 | 1, nowl, nowr);
    return sum;
}

int query_tree(int p)
{
    return query(1, n, 1, rk[p], rk[p] + sz[p] - 1);
}

void modify_tree(int p, int value)
{
    modify(1, n, 1, rk[p], rk[p] + sz[p] - 1, 1ll << value);
}
signed main()
{
    cin >> n >> q;
    for (int i = 1; i <= n; i++)
        cin >> col[i];
    for (int i = 1; i < n; i++)
    {
        int u, v;
        cin >> u >> v;
        z[u].push_back(v);
        z[v].push_back(u);
    }
    dfs1(1);
    dfs2(1, 0);
    build(1, n, 1);
    while (q--)
    {
        int o;
        cin >> o;
        if (o == 1)
        {
            int u, c;
            cin >> u >> c;
            modify_tree(u, c);
        }
        else
        {
            int u;
            cin >> u;
            cout << __builtin_popcountll(query_tree(u)) << '\n';
        }
    }
    return 0;
}