CF620E--New Year Tree(线段树)

Posted 2020-10-10 zQ

CodeForces - 620E

1.题意

　　　　对于一棵树进行以下两种操作:

                 1  v c　将根为v的子树全部染色为c

 　　　　  2 v 查询根为v的子树的颜色的个数

2.思路

 　　　　首先利用dfs时间序将树转化为线性结构,然后利用线段树进行区间颜色数的处理,然后因为只有60种颜色,那么我们可以用一个long long 值来保存一个去区间颜色数(也就是对于颜色i,那么对应的longlong值的二进制位为1,一个区间的颜色数也就是区间值的二进制的1的个数),然后区间合并则直接对两个区间的值进行或运算(|);

　　　　需要注意的是移位时防止溢出.

3.代码

#define debug printf("%s: %d\n", __FUNCTION__, __LINE__);

#include <bits/stdc++.h>
using namespace std;
const int N = 4e5+5;
typedef long long ll;
#define ls rt<<1
#define rs rt<<1|1

// 二进制记录数居

struct Edge{
    int to, next;
};

int head[N];
Edge e[N<<1];
int tot;

void addEdge(int u, int v) {
    e[++ tot] = Edge{u, head[v]};
    head[v] = tot;
    e[++ tot] = Edge{v, head[u]};
    head[u] = tot;
}

int l[N], r[N];
int dfs_clock;
int pos[N];

// dfs将树转化为线性结构
void dfs(int u, int fa) {
    l[u] = ++ dfs_clock;
    pos[dfs_clock] = u;
    for(int i = head[u]; i; i = e[i].next) {
        int to = e[i].to;
        if(to == fa) continue;
        dfs(to, u);
    }
    r[u] = dfs_clock;    
}

int c[N];
ll v[N<<2];
int flag[N<<2];

inline void pushUp(int rt) {
    v[rt] = v[ls]|v[rs];
}

void pushDown(int rt) {
    if(flag[rt]) {
        flag[ls] = flag[rs] = flag[rt];
        v[ls] = (1ll<<flag[rt]);
        v[rs] = (1ll<<flag[rt]);   // 注意防止溢出
        flag[rt] = 0;
    }
}

void build(int rt, int l, int r) {
    if(l == r) {
        v[rt] = (1ll<<c[pos[l]]);
        return ;
    }
    int mid = (l + r) >> 1;
    build(ls, l, mid);
    build(rs, mid+1, r); 
    pushUp(rt);
}

void update(int rt, int l, int r, int a, int b, int ci) {
    if(a <= l && r <= b) {
        flag[rt] = ci;
        v[rt] = (1ll<<ci);
        return ;
    } 
    pushDown(rt);
    int mid = (l + r) >> 1;
    if(a <= mid) {
        update(ls, l, mid, a, b, ci);
    }
    if(b > mid) {
        update(rs, mid+1, r, a, b, ci);
    }
    pushUp(rt);
}

ll query(int rt, int l, int r, int a, int b) {
    if(a <= l && r <= b) {
        return v[rt];
    } 
    pushDown(rt);
    ll res = 0;
    int mid = (l + r) >> 1;
    if(a <= mid) {
        res |= query(ls, l, mid, a, b);
    }
    if(b > mid) {
        res |= query(rs, mid+1, r, a, b);
    }
    pushUp(rt);
    return res;
}

int main() {
    int n, m;
    scanf("%d%d", &n, &m);
    for(int i = 1; i <= n; ++ i) {
        scanf("%d", &c[i]);
    }     
    for(int i = 1; i < n; ++ i) {
        int x, y;
        scanf("%d%d", &x, &y);
        addEdge(x, y);
    }

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

    while(m --) {
        int k;
        scanf("%d", &k);
        if(k == 1) {
            int x, y;
            scanf("%d%d", &x, &y);
            update(1, 1, n, l[x], r[x], y);  
        }
        else {
            int x;
            scanf("%d", &x);
            ll res = query(1, 1, n, l[x], r[x]);
            int s = 0;
            while(res) {
                s += res&1;
                res >>= 1;
            }
            printf("%d\n", s);
        }
    }
    return 0;
}