E. New Year Tree(树剖+bitset)

135 阅读3分钟

题目

The New Year holidays are over, but Resha doesn't want to throw away the New Year tree. He invited his best friends Kerim and Gural to help him to redecorate the New Year tree.

The New Year tree is an undirected tree with n vertices and root in the vertex 1.

You should process the queries of the two types:

Change the colours of all vertices in the subtree of the vertex v to the colour c. Find the number of different colours in the subtree of the vertex v.

输入

The first line contains two integers n, m (1 ≤ n, m ≤ 4·105) — the number of vertices in the tree and the number of the queries.

The second line contains n integers ci (1 ≤ ci ≤ 60) — the colour of the i-th vertex.

Each of the next n - 1 lines contains two integers xj, yj (1 ≤ xj, yj ≤ n) — the vertices of the j-th edge. It is guaranteed that you are given correct undirected tree.

The last m lines contains the description of the queries. Each description starts with the integer tk (1 ≤ tk ≤ 2) — the type of the k-th query. For the queries of the first type then follows two integers vk, ck (1 ≤ vk ≤ n, 1 ≤ ck ≤ 60) — the number of the vertex whose subtree will be recoloured with the colour ck. For the queries of the second type then follows integer vk (1 ≤ vk ≤ n) — the number of the vertex for which subtree you should find the number of different colours.

解法

树链剖分维护比特位

Code

#include <iostream>
#include <cstring>
#include <algorithm>
#include <map>
#include <queue>
#include <vector>
#include <cmath>
#include <set>
#include <stack>
#include <bits/stdc++.h>
using namespace std;
using ll = long long;
#define int long long
#define endl '\n'
#define pb push_back
#define NO cout << "NO" << endl;
#define YES cout << "YES" << endl;
#define fi first
#define se second
#define all(x) (x).begin(),(x).end()
#define rep(i,a,n) for (int i=a;i<=n;i++)
#define per(i,a,n) for (int i=n;i>=a;i--)
typedef vector<int> VI;
typedef pair<int,int> PII;
typedef vector<VI> VII;
ll MOD = 998244353;
ll powmod(ll a,ll b) {ll res=1;a%=MOD; assert(b>=0); for(;b;b>>=1){if(b&1)res=res*a%MOD;a=a*a%MOD;}return res;}
mt19937 mrand(random_device{}()); 
const int N = 4e5 + 10;
int w[N];
int n,m;
int son[N],id[N],fa[N],cnt,dep[N],siz[N],top[N],wt[N]; 
vector<int> e[N];
int f[N << 2],Lazy[N << 2];
void Built(int k,int l,int r) {
    if( l == r) {
        f[k] = (1ll << wt[l]);
        // cout << f[k] << endl;
        return;
    }
    int m = l + r >> 1;
    Built(k<<1,l,m);
    Built(k<<1|1,m+1,r);
    f[k] = f[k << 1] | f[k << 1 | 1];
}
void update(int k,int l,int r,int L,int R,int s)
{
    if( l == L && r == R)
    {
        Lazy[k] = 1ll << s;
        f[k] = 1ll << s;
        return;
    }
    if(Lazy[k]) {
        f[k << 1] = Lazy[k];
        f[k << 1|1] = Lazy[k];
        Lazy[k <<1] = Lazy[k];
        Lazy[k << 1|1] = Lazy[k];
        Lazy[k] = 0;

    }
    int m = (l + r) >> 1;
    if(R <= m) update(k << 1,l,m,L,R,s);
    else
    {
       if(L > m) update(k << 1|1,m+1,r,L,R,s);
       else update(k << 1,l,m,L,m,s),update(k <<1|1,m+1,r,m+1,R,s);
    }
    f[k] = f[k << 1] | f[k << 1 | 1];
}
int calc(int k,int l,int r,int x,int y)
{
    if(l == x && r == y){
        return f[k];
    }
    if(Lazy[k]) {
        f[k << 1] = Lazy[k];
        f[k << 1|1] = Lazy[k];
        Lazy[k <<1] = Lazy[k];
        Lazy[k << 1|1] = Lazy[k];
        Lazy[k] = 0;

    }
    int res = 0;
    int m = (l + r) >> 1;
    if(y <= m) res |= calc(k << 1,l,m,x,y);
    else {
        if(x > m) res |= calc(k<<1|1,m+1,r,x,y);
        else res |= (calc(k<<1,l,m,x,m) | calc(k<<1|1,m+1,r,m+1,y));
    }
    f[k] = f[k << 1] | f[k << 1 | 1];
    return res;
}
void dfs1(int x,int f,int deep) { // 找到每个结点的子节点数量和重儿子
   dep[x] = deep;
   fa[x] = f;
   siz[x] = 1;
   int maxson = -1;
   for(auto u : e[x]) {
     if(u == f) continue;
     dfs1(u,x,deep + 1);
     siz[x] += siz[u];
     if(siz[u] > maxson) {
           son[x] = u;
           maxson = siz[u];
     }
   }
}
void dfs2(int x,int topf) {
    id[x] = ++cnt;
    wt[cnt] = w[x];
    top[x] = topf;
    if(!son[x]) return;
    dfs2(son[x],topf);
    for(auto u : e[x]) {
        if(u == fa[x] || u == son[x]) continue;
        dfs2(u,u);
    }
}
int qRange(int x,int y){ // x -> y 区间求和
    int ans=0;
    while(top[x] != top[y]){//当两个点不在同一条链上 
        if(dep[top[x]] < dep[top[y]])swap(x,y);//把x点改为所在链顶端的深度更深的那个点
        ans |= calc(1,1,n,id[top[x]],id[x]);//ans加上x点到x所在链顶端 这一段区间的点权和//按题意取模 
        x=fa[top[x]];//把x跳到x所在链顶端的那个点的上面一个点
    }
    //直到两个点处于一条链上
    if(dep[x] > dep[y])swap(x,y);//把x点深度更深的那个点
    int res = 0;
    ans |= calc(1,1,n,id[x],id[y]);//这时再加上此时两个点的区间和即可
    return ans;
}
void updRange(int x,int y,int k){// 区间修改
    while(top[x] != top[y]){
        if(dep[top[x]] < dep[top[y]]) swap(x,y);
        update(1,1,n,id[top[x]],id[x],k);
        x = fa[top[x]];
    }
    if(dep[x] > dep[y])swap(x,y);
    update(1,1,n,id[x],id[y],k);
}
int qSon(int x){ //单点查询
    int res = 0;
    return calc(1,1,n,id[x],id[x] + siz[x]-1);//子树区间右端点为id[x]+siz[x]-1 
     // return res;
}
void updSon(int x,int k) { // 更新子树
    update(1,1,n,id[x],id[x] + siz[x]-1,k);//
    // update(1,1,n,id[x],id[x],k); // 更新单个节点

}
// r是根节点
// dfs1(r,0,1), dfs(r,r);
// 树剖边权化点
void solve()
{
    cin >> n >> m;
    rep (i, 1, n) cin >> w[i];
    rep (i, 1, n - 1) {
        int u, v; cin >> u >> v;
        e[u].pb(v);
        e[v].pb(u);
    }
    int r = 1;
    dfs1(r,0,1), dfs2(r,r);
    Built(1, 1, n);
    while (m --) {
        int op; cin >> op;
        if (op == 1) {
            int v, c; cin >> v >> c;
            updSon(v, c);
        } else {
            int x; cin >> x;
            ll res = 0;
            res = qSon(x);
            
            // cout << x << "---" << res << endl;
            ll ans = 0;
            while (res) {
                if (res & 1) ans ++;
                res /= 2;
            }
            cout << ans << endl;
            // cout << 1 << endl;
        }
    }
    // cout << (4 | 16) << endl;
}
signed main()
{
    ios::sync_with_stdio(false);
    cin.tie(0); cout.tie(0);
    // int T;cin >> T;
    // while ( T -- )
    solve();
    return 0;
}