Codeforces 620E New Year Tree dfs序+线段树+状态压缩

My Solution

题意：给定一棵树，每个节点都有颜色，然后询问子树上有多少种不同的颜色。

dfs序+线段树+状态压缩
由于只有60种颜色（2^60 < 2^63），所以可以直接用二进制压位。

即sum[Ind]维护的是该区间的一个状态，从右向左第i位表示第i种颜色在该区间是否出现。

然后用上线段树区间修改+区间查询即可。

时间复杂度 O(nlogn)

空间复杂度 O(4*n)

#include 
#include 
#include 
#include 
using namespace std;
typedef long long LL;
const int MAXN = 4e5 + 8;

vector sons[MAXN];
int color[MAXN];
//dfs序
int p1[MAXN], p2[MAXN], ti = 0;
int dfsnum[MAXN];  //这个按情况是否需要。
inline void get_dfs_list(int u, int fa){
    p1[u] = ++ti;
    dfsnum[ti] = u; //
    int sz = sons[u].size(), i, v;
    for(i = 0; i < sz; i++){
        v = sons[u][i];
        if(v == fa) continue;
        get_dfs_list(v, u);
    }
    p2[u] = ti;
}
//线段树
LL sum[4*MAXN], lazy[4*MAXN];
int size;
inline void pushup(int Ind){
    sum[Ind] = sum[Ind<<1] | sum[(Ind<<1)|1];
}

inline void pushdown(int Ind){
    sum[Ind<<1] = lazy[Ind];
    sum[(Ind<<1)|1] = lazy[Ind];
    lazy[Ind<<1] = lazy[Ind];
    lazy[(Ind<<1)|1] = lazy[Ind];
    lazy[Ind] = 0;
}

inline LL _Query(int a, int b, int l, int r, int Ind){
    if(a <= l && b >= r) return sum[Ind];
    int mid = (l+r)>>1;
    if(lazy[Ind]) pushdown(Ind);
    LL ret = 0;
    if(a <= mid) ret |= _Query(a, b, l, mid, Ind<<1); if(b > mid) ret |= _Query(a, b, mid + 1, r, (Ind<<1)|1);
    //pushup(Ind);
    return ret;
}

inline void _Modify(int a, int b, int l, int r, int Ind, LL d){
     if(a <= l && b >= r){
        sum[Ind] = d;
        lazy[Ind] = d;
        return;
    }
    int mid = (l+r)>>1;
    if(lazy[Ind]) pushdown(Ind);
    if(a <= mid) _Modify(a, b, l, mid, Ind<<1, d); if(b > mid) _Modify(a, b, mid + 1, r, (Ind<<1)|1, d);
    pushup(Ind);
}

inline void _build(int l, int r, int Ind){
    if(l == r){
        sum[Ind] = 1ll << (color[dfsnum[l]] - 1); return; } int mid = (l+r)>>1;
    _build(l, mid, Ind<<1);
    _build(mid + 1, r, (Ind<<1)|1); pushup(Ind); } inline LL Query(int a, int b) {return _Query(a, b, 1, size, 1);} inline void Modify(int a, int b, LL d){return _Modify(a, b, 1, size, 1, d);} int main() { #ifdef LOCAL freopen("c.txt", "r", stdin); //freopen("c.out", "w", stdout); int T = 3; while(T--){ #endif // LOCAL ios::sync_with_stdio(false); cin.tie(0); int n, m, i, u, v, root, type, xi, ci, ans; LL val; cin >> n >> m;
    for(i = 1; i <= n; i++){ cin >> color[i];
    }
    for(i = 1; i < n; i++){ cin >> u >> v;
        sons[u].push_back(v);
        sons[v].push_back(u);
    }
    root = 1;
    ti = 0;
    get_dfs_list(root, -1);
    size = ti;
    _build(1, size, 1);
    while(m--){
        cin >> type;
        if(type == 1){
            cin >> xi >> ci;
            Modify(p1[xi], p2[xi], 1ll <<(ci-1)); } else{ cin >> xi;
            val = Query(p1[xi], p2[xi]);
            ans = 0;
            while(val){
                if(val & 1){
                    ans++;
                }
                val >>= 1;
            }
            cout << ans << "n";
        }
    }


    #ifdef LOCAL
    for(i = 1; i <= n; i++){
        sons[i].clear();
    }
    memset(sum, 0, sizeof sum);
    memset(lazy, 0, sizeof lazy);
    cout << endl;
    }
    #endif // LOCAL
    return 0;
}

非特殊说明，本博所有文章均为博主原创，未经许可不得转载。

https://www.prolightsfxjh.com/

Thank you!

                                                                                                                                             ------from ProLightsfx