51nod 1819 黑白树V2（树链剖分）

第一次写如此复杂的树链剖分，

感觉自己代码能力还是挺不错的，没有调试太久（2个小时）

最后代码量高达11K orz（大部分都是重复的线段树代码，以后可以考虑优化一下代码量）

题解：

首先就是要进行一次树链剖分，

然后线段树维护的内容是比较复杂的，核心是这样的

线段树上的一个结点维护以下信息

v[2][2] 第一维代表deep%2，第二维代表是黑还是白（黑是1），v[0][1]意思就是深度为偶数、黑色结点的答案

　　　　这里答案是指的维护所有轻儿子子树的答案

num[2][2] 维度代表和上面一样，维护的是所有轻儿子子树中满足该颜色、该深度的结点总数

tag[2] 维度代表deep%2，为1就是代表与当前颜色相反，为0就是相同

ftag[2] 维度依然代表deep%2，永久化标记，维护一个区间被打的永久化标记的情况

这里简单说一下维护的细节和思路吧

1、对于第一个操作修改点的颜色

实际上就是先找到这个点，然后直接修改它的颜色（不打标记），接下来就是不断往上跳，去修改fa[top[x]]的值，因为修改单点颜色会影响它所维护的答案

2、对于第二个操作

也是首先找到这个点，然后给它先打上一个永久化标记，并对它所在的重链打一个懒惰标记。

然后需要去依次修改fa[top[x]]的值，这时我们需要知道它实际改变了多少，所以我们就从底向上统计所有的永久化标记

这样就可以知道实际改变了多少，然后每修改一条重链的点，就删去这条重链上永久化标记的影响，就可以了

3、对于第三个操作

统计答案，要按链来进行

一开始对于第一条重链，首先统计该点x下有多少个结点，得到这个n以后乘x即可，然后在直接在线段树上查询top[x] ~fa[x]这个区间的值，加上即可

然后跳到下一条重链，也是统计新的x有多少个结点，得到这个n以后需要减去son[x]的相应结点个数，再乘x，接下来还是查询top[x]~fa[x]，依次进行

注意在跳跃的过程中，也要不断修改永久化标记的影响，这样才会得到真正的答案。

4、对于线段树初始值

要进行一个简单的树型dp，这个比较水就不说了。

写代码的时候也发现到一些问题，比如说最好写赋值型修改，不要返回一个数组指针（会因为局部变量丢失而崩掉）

然后这个代码重复率高的问题也需要解决。

#include <iostream>
#include <cstdio>
#include <cstring>
#include <vector>
#include <hash_map>
#define fi first
#define se second
#define pb push_back
using namespace std;
using namespace __gnu_cxx;
typedef pair<int, int> PII;
typedef long long LL;
const int maxn = 200005;
int deep[maxn], fa[maxn], sz[maxn], son[maxn], top[maxn], down[maxn], c[maxn];
PII dp[maxn][2], dps[maxn][2];
int tot, n;
vector<int>G[maxn];
hash_map<int, int> H, iH;

void print(PII *A){
    cout<<A[0].fi<<" "<<A[0].se<<endl;
    cout<<A[1].fi<<" "<<A[1].se<<endl;
    cout<<endl;
}
void print(int *A){
    cout<<A[0]<<" "<<A[1]<<endl; cout<<endl;
}

int dfs1(int x){
    deep[x] = deep[fa[x]]+1;
    sz[x] = 1;
    for(auto to : G[x]){
        if(to == fa[x]) continue;
        fa[to] = x; sz[x] += dfs1(to);
        son[x] = sz[to] > sz[son[x]] ? to : son[x];
    }
    return sz[x];
}
void dfs2(int x){
    iH[H[x] = ++tot] = x;
    top[x] = son[fa[x]] == x ? top[fa[x]] : x;
    if(son[x]) dfs2(son[x]); down[x] = down[son[x]];
    for(auto to : G[x]){
        if(to == fa[x] || to == son[x]) continue;
        dfs2(to);
    }
    if(!son[x]) down[x] = x;
}
void dfs3(int x){
    int d = deep[x]%2;
    if(c[x]) dp[x][d].se = 1; else dp[x][d].fi = 1;
    if(c[x]) dps[x][d].se = 1; else dps[x][d].fi = 1;
    for(auto to : G[x]){
        if(to == fa[x]) continue;
        dfs3(to);
        for(int i = 0; i < 2; i++){
            dps[x][i].fi += dps[to][i].fi;
            dps[x][i].se += dps[to][i].se;
        }
        if(to == son[x]) continue;
        for(int i = 0; i < 2; i++) {
            dp[x][i].fi += dps[to][i].fi;
            dp[x][i].se += dps[to][i].se;
        }
    }
}

void Mulnum(PII D[2], PII A[2], PII B[2]){
    PII C[2];
    for(int i = 0; i < 2; i++) C[i].fi = A[i].fi + B[i].fi, C[i].se = A[i].se + B[i].se;
    for(int i = 0; i < 2; i++) D[i].fi = C[i].fi, D[i].se = C[i].se;
}
void Multag(int *D, int *A, int *B){
    int C[2];
    for(int i = 0; i < 2; i++) C[i] = A[i]^B[i];
    for(int i = 0; i < 2; i++) D[i] = C[i];
}
void MulNF(PII *A, int *B){
    for(int i = 0; i < 2; i++) if(B[i]) swap(A[i].fi, A[i].se);
}
LL MulVF(PII *A, int *B){
    LL ans = 0;
    for(int i = 0; i < 2; i++) ans += B[i] ? A[i].fi : A[i].se;
    return ans;
}
void Sub(PII D[2], PII A[2], PII B[2]){
    PII C[2];
    for(int i = 0; i < 2; i++) C[i].fi = A[i].fi - B[i].fi, C[i].se = A[i].se - B[i].se;
    for(int i = 0; i < 2; i++) D[i].fi = C[i].fi, D[i].se = C[i].se;
}

struct Tree{
    PII num[2], v[2];
    int tag[2];
    int ftag[2];
}tree[maxn*4];
void Puttag(int o, int* tag){
    for(int i = 0; i < 2; i++) if(tag[i]) swap(tree[o].num[i].fi, tree[o].num[i].se);
    for(int i = 0; i < 2; i++) if(tag[i]) swap(tree[o].v[i].fi, tree[o].v[i].se);
    Multag(tree[o].tag, tree[o].tag, tag);
}
void Pushdown(int o){
    if(tree[o].tag[0] == 0 && tree[o].tag[1] == 0) return;
    Puttag(o*2, tree[o].tag);
    Puttag(o*2+1, tree[o].tag);
    tree[o].tag[0] = tree[o].tag[1] = 0;
}
void Maintain(int o){
    Mulnum(tree[o].num, tree[o*2].num, tree[o*2+1].num);
    Mulnum(tree[o].v, tree[o*2].v, tree[o*2+1].v);
    Multag(tree[o].ftag, tree[o*2].ftag, tree[o*2+1].ftag);
}

void Query(int o, int l, int r, int k){
    if(l == r){
        print(tree[o].v);
        print(tree[o].num);
        cout<<"******"<<endl;
        return;
    }
    int mid = (l+r)>>1; Pushdown(o);
    if(k <= mid) Query(o*2, l, mid, k);
    else Query(o*2+1, mid+1, r, k);
    Maintain(o);
}

void colchange(int o, int v, PII num[2]){
    Mulnum(tree[o].num, tree[o].num, num);
    for(int i = 0; i < 2; i++){
        tree[o].v[i].se = v*tree[o].num[i].se;
        tree[o].v[i].fi = v*tree[o].num[i].fi;
    }
}

int _Schange(int o, int l, int r, int k){
    if(l == r){
        int d = deep[iH[k]]%2, col = c[iH[k]]; c[iH[k]] ^= 1;
        col ^= tree[o].tag[d];
        PII temp[2]; temp[0] = temp[1] = {0, 0};
        if(col) temp[d] = {1, -1}; else temp[d] = {-1, 1};
        colchange(o, iH[k], temp);
        return col^1;
    }
    int mid = (l+r)>>1, col; Pushdown(o);
    if(k <= mid) col = _Schange(o*2, l, mid, k);
    else col = _Schange(o*2+1, mid+1, r, k);
    Maintain(o);
    return col;
}

void _ftagQuery(int o, int l, int r, int L, int R, int *ans){
    if(L <= l && r <= R) {
        Multag(ans, ans, tree[o].ftag);
        return;
    }
    int temp[2] = {0, 0};
    int mid = (l+r)>>1; Pushdown(o);
    if(L <= mid) _ftagQuery(o*2, l, mid, L, R, temp);
    if(R > mid) _ftagQuery(o*2+1, mid+1, r, L, R, temp);
    Maintain(o);
    Multag(ans, ans, temp);
}

void _Coladd(int o, int l, int r, int k, PII* num){
    if(l == r){
        colchange(o, iH[k], num);
        return;
    }
    int mid = (l+r)>>1; Pushdown(o);
    if(k <= mid) _Coladd(o*2, l, mid, k, num);
    else _Coladd(o*2+1, mid+1, r, k, num);
    Maintain(o);
}
void _numQuery(int o, int l, int r, int L, int R, PII* ans){
    if(L <= l && r <= R) {
        Mulnum(ans, ans, tree[o].num);
        return;
    }
    int mid = (l+r)>>1; Pushdown(o);
    PII temp[2] = {{0, 0}, {0, 0}};
    if(L <= mid) _numQuery(o*2, l, mid, L, R, temp);
    if(R > mid)  _numQuery(o*2+1, mid+1, r, L, R, temp);
    Maintain(o);
    Mulnum(ans, ans, temp);
}
void _vQuery(int o, int l, int r, int L, int R, PII *ans){
    if(L <= l && r <= R) {
        Mulnum(ans, ans, tree[o].v);
        return;
    }
    int mid = (l+r)>>1; Pushdown(o);
    PII temp[2] = {{0, 0}, {0, 0}};
    if(L <= mid) _vQuery(o*2, l, mid, L, R, temp);
    if(R > mid) _vQuery(o*2+1, mid+1, r, L, R, temp);
    /*cout<<l<<" "<<r<<"-"<<endl;
    print(ans);
    cout<<"---------"<<endl;*/
    Maintain(o);
    Mulnum(ans, ans, temp);
}
void _ftagchange(int o, int l, int r, int k, int d){
    if(l == r){
        tree[o].ftag[d] ^= 1;
        return;
    }
    int mid = (l+r)>>1; Pushdown(o);
    if(k <= mid) _ftagchange(o*2, l, mid, k, d);
    else _ftagchange(o*2+1, mid+1, r, k, d);
    Maintain(o);
}
void _tagchange(int o, int l, int r, int L, int R, int d){
    if(L <= l && r <= R){
        int temp[2] = {0, 0}; temp[d] = 1;
        Puttag(o, temp);
        return;
    }
    int mid = (l+r)>>1; Pushdown(o);
    if(L <= mid) _tagchange(o*2, l, mid, L, R, d);
    if(R > mid) _tagchange(o*2+1, mid+1, r, L, R, d);
    Maintain(o);
}



int Schange(int o, int l, int r, int k){ return _Schange(o, l, r, H[k]); }
void ftagQuery(int o, int l, int r, int L, int R, int* ans) { _ftagQuery(o, l, r, H[L], H[R], ans);}
void Col_add(int o, int l, int r, int x, PII* num) { _Coladd(o, l, r, H[x], num); }
void numQuery(int o, int l, int r, int L, int R, PII* ans) { _numQuery(o, l, r, H[L], H[R], ans); }
void vQuery(int o, int l, int r, int L, int R, PII* ans) { _vQuery(o, l, r, H[L], H[R], ans); }
void ftagchange(int o, int l, int r, int x, int d) { _ftagchange(o, l, r, H[x], d); }
void tagchange(int o, int l, int r, int L, int R, int d) { _tagchange(o, l, r, H[L], H[R], d);}

void T_Insert(int x){
    int y = top[x];
    int col = Schange(1, 1, n, x), d = deep[x]%2;
    int ftag[2] = {0, 0};
    while(y != 1){
        x = fa[y];
        y = top[x];
        ftagQuery(1, 1, n, y, x, ftag);
        if(ftag[d]) col ^= 1;
        PII temp[2] = {{0, 0}, {0, 0}};
        if(col) temp[d] = {-1, 1}; else temp[d] = {1, -1};
        Col_add(1, 1, n, x, temp);
    }
}

void find_ftag(int x, int *ftag){
    ftag[0] = ftag[1] = 0;
    if(x == 0) return;
    int y = top[x];
    ftagQuery(1, 1, n, y, x, ftag);
    while(y != 1){
        x = fa[y];
        y = top[x];
        int temp[2] = {0, 0};
        ftagQuery(1, 1, n, y, x, temp);
        Multag(ftag, ftag, temp);
    }
}

LL T_Query(int x){
    int ftag[2] = {0, 0};
    find_ftag(fa[top[x]], ftag);
    LL ans = 0;
    PII last[2] = {{0, 0}, {0, 0}};
    //for(int i = 1; i <= 3; i++) Query(1, 1, n, i);
    while(1){
        PII num[2] = {{0, 0}, {0, 0}};
        numQuery(1, 1, n, x, down[x], num);
        //print(num);
        MulNF(num, ftag);

        Sub(num, num, last);
        last[0] = last[1] = {0, 0};
        ans += (num[0].se + num[1].se)*x;
        //cout<<ans<<"*"<<endl;
        if(x != top[x]){
            PII v[2] = {{0, 0}, {0, 0}};
            vQuery(1, 1, n, top[x], fa[x], v);
            ans += MulVF(v, ftag);
        }
        //cout<<ans<<"*"<<endl;
        numQuery(1, 1, n, top[x], down[x], last);
        MulNF(last, ftag);
        x = fa[top[x]];
        if(x == 0) break;

        ftagQuery(1, 1, n, top[x], x, ftag);
    }
    return ans;
}

void T_Change(int x){
    PII num1[2] = {{0, 0}, {0, 0}}, num2[2] = {{0, 0}, {0, 0}};
    //cout<<x<<" "<<down[x]<<endl;
    numQuery(1, 1, n, x, down[x], num1);
    //print(num1);
    ftagchange(1, 1, n, x, (deep[x]+1)%2);
    tagchange(1, 1, n, x, down[x], (deep[x]+1)%2);
    numQuery(1, 1, n, x, down[x], num2);
    //print(num2);
    int y = top[x];
    while(y != 1){
        x = fa[y];
        y = top[x];
        int ftag[2] = {0, 0};
        ftagQuery(1, 1, n, y, x, ftag);
        MulNF(num1, ftag); MulNF(num2, ftag);
        PII temp[2];
        for(int i = 0; i < 2; i++) {
            temp[i].fi = num2[i].fi - num1[i].fi;
            temp[i].se = num2[i].se - num1[i].se;
        }
        Col_add(1, 1, n, x, temp);
    }
}

void Insert(int o, int l, int r, int k, PII *num, LL v){
    if(l == r){
        for(int i = 0; i < 2; i++) {
            tree[o].num[i].fi = num[i].fi;
            tree[o].v[i].fi = num[i].fi*v;
            tree[o].num[i].se = num[i].se;
            tree[o].v[i].se = num[i].se*v;
        }
        return;
    }
    int mid = (l+r)>>1;
    if(k <= mid) Insert(o*2, l, mid, k, num, v);
    else Insert(o*2+1, mid+1, r, k, num, v);
    Maintain(o);
}

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

    fa[1] = 0;
    dfs1(1); dfs2(1); dfs3(1);
    for(int i = 1; i <= n; i++) Insert(1, 1, n, H[i], dp[i], i);
    //for(int i = 1; i <= n; i++) Query(1, 1, n, i);
    while(q--){
        scanf("%d %d", &x, &y);
        if(x == 1){
            T_Change(y);
        } else if(x == 2){
            T_Insert(y);
            //Query(1, 1, n, H[5]);
        } else if(x == 3){
            printf("%lld
", T_Query(y));
        }
    }
    return 0;
}