C - New Year Tree CodeForces - 620E（dfs 序+线段树（维护颜色用状态压缩））

题意

1. 给你一颗 n 个节点的树，1 为根节点，每个节点被染上了颜色，然后就是 m 次查询。 查询的方式有两种 1，将以 z 为根的子树的结点全部更新为颜色 X 2，问以 z 为根的子树的结点的不同颜色数量。

思路

1. 先用 dfs 序，给这棵树上的节点重新变一下编号（盖上时间戳）， 这样把树形问题变成了 线性区间问题，就可以用线段树及决问题了，
2. 由于颜色种类少，所以用线段树维护颜色的话可以用二进制状态压缩（每个 bit 位表示一种颜色），来表示，这样当两个颜色合并的时候，只需要把 两种颜色状态相或 | 就行了（这样操作方便好多）。

代码

#include <iostream>
#include <stdio.h>
#include <string.h>
#include <math.h>
#include <algorithm>
#include <string>
#include <queue>
#include <map>
#include <bitset>
#include <vector>
using namespace std;
void fre() { system("clear"), freopen("A.txt", "r", stdin); freopen("Ans.txt","w",stdout); }
void Fre() { system("clear"), freopen("A.txt", "r", stdin);}
void Run(int x = 0) {     
#ifdef ACM  //宏定义免注释 freopen
    if (! x) fre(); else Fre();
#endif
}
#define ios ios::sync_with_stdio(false)
#define Pi acos(-1)
#define pb push_back
#define fi first
#define se second
#define db double
#define ll long long
#define ull unsigned long long
#define Pir pair<ll, ll>
#define m_p make_pair
#define for_(i, s, e) for(ll i = (ll)(s); i <= (ll)(e); i ++)
#define rep_(i, e, s) for(ll i = (ll)(e); i >= (ll)(s); i --)
#define memset(a, b, c) memset(a, (int)b, c);
#define size() size() * 1LL
#define sc scanf
#define pr printf
#define sd(a) sc("%lld", &a)
#define ss(a) sc("%s", a)
#define __  pr( "------------------------\n" );
#define ___ pr("\n------------------------\n");
#define inf 0x3f3f3f3f
#define INF 0x3f3f3f3f3f3f3f3f
#define esp 1e-7
#define mod (ll)(1e9 + 7)
/*=========================ACMer===========================*/
const int mxn = 4e5 + 10;
vector<int> e[mxn];
int a[mxn];
ll color[mxn];              //这个不要忘记开 longlong
int L[mxn], R[mxn], tim;

void dfs(int u, int p)
{
    L[u] = ++ tim;
    color[tim] = (1LL << a[u]);
    for_(i, 0, e[u].size() - 1)
    {
        int v = e[u][i];
        if(v == p) continue;
        dfs(v, u);
    }
    R[u] = tim;
}

struct Tree
{
    int l, r, lazy;
    ll state;
} tree[mxn << 2];

void push_up(int rt)
{
    tree[rt].state = tree[rt << 1].state | tree[rt << 1 | 1].state;
}

void push_down(int rt)
{
    if(tree[rt].lazy)
    {
        tree[rt << 1].lazy = tree[rt << 1 | 1].lazy = 1;
        tree[rt << 1].state = tree[rt << 1 | 1].state = tree[rt].state;
        tree[rt].lazy = 0;
    }
}

void build(int rt, int l, int r)
{
    tree[rt].l = l, tree[rt].r = r, tree[rt].lazy = 0;
    if(l == r)
    {
        tree[rt].state = color[l];         //color[l] 的值在 dfs 的时候被赋值了
        return;
    }

    int md = (l + r) >> 1;
    build(rt << 1, l, md);
    build(rt << 1 | 1, md + 1, r);
    push_up(rt);
}

void update(int rt, int l, int r, int k)
{
    if(tree[rt].l >= l && tree[rt].r <= r)
    {
        tree[rt].state = (1LL << k);
        tree[rt].lazy = 1;
        return;
    }

    push_down(rt);
    if(tree[rt << 1].r >= l) update(rt << 1, l, r, k);
    if(tree[rt << 1 | 1].l <= r) update(rt << 1 | 1, l, r, k);
    push_up(rt);
}

ll query(int rt, int l, int r)
{
    if(tree[rt].l >= l && tree[rt].r <= r)
    {
        return tree[rt].state;
    }

    push_down(rt);

    ll sum = 0;
    if(tree[rt << 1].r >= l) sum |= query(rt << 1, l, r);
    if(tree[rt << 1 | 1].l <= r) sum |= query(rt << 1 | 1, l , r);
    return sum;
}

void init(ll n)
{
    for_(i, 1, n) e[i].clear();
    tim = 0;
}

int main()
{
    Run();
    int n, m;
    while(sc("%d %d", &n, &m) != EOF)
    {
        init(n);
        for_(i, 1, n) sc("%d", &a[i]);
        int u, v;
        for_(i, 1, n - 1)
        {
            sc("%d %d", &u, &v);
            e[u].pb(v);
            e[v].pb(u);
        }

        dfs(1, -1);
        build(1, 1, n);

        int op, x, y;
        while(m --)
        {
            sc("%d", &op);
            if(op == 1)
            {
                sc("%d %d", &x, &y);
                update(1, L[x], R[x], y);
            }
            else
            {
                sc("%d", &x);
                ll state = query(1, L[x], R[x]);
                int ans = 0;
                while(state)
                {
                    if(state & 1LL)
            
                        ans ++;
                    
                    state >>= 1;
                }
                pr("%d\n", ans);
            }
        }
    }

    return 0;
}