bzoj 3637: Query on a tree VI （树链剖分+树状数组）

题目描述

传送门

题目大意：给出一棵树，m个操作，每次操作要么更改某个点的颜色要么询问与某个点颜色相同的联通块大小。

题解

对于每个点维护$T[0/1][x]$表示x的子树中当x的颜色为0/1时连通块的大小。
那么每次查询的时候我们只需要找到与x颜色相同的深度最浅的节点（x到该点路径上的所有节点必须与x同色），然后查询这个点子树中col[x]连通块的大小
对于修改，注意我们在维护的时候都是不包括x本身的，所以修改的时候只需要修改他上面的一串节点。
如果fa[x]与x的颜色不同，那么我们要在$T[col[x]][fa[x]]$中减掉x的贡献。设t为fa[x]走相同颜色所能到达的最浅节点，需要在fa[x]->fa[t]这条链中加入x的贡献。注意这里的贡献不只是x本身，还包括x子树的信息。
如果fa[x]与x的颜色相同，依然需要在$T[col[x]^1][fa[x]]$中加入x的贡献。设t为fa[x]走相同颜色所能到达的最浅节点，需要在fa[x]->fa[t]这条链中减掉x的贡献。

代码

#include<iostream>
#include<cstdio>
#include<algorithm>
#include<cmath>
#include<cstring>
#define N 200033
using namespace std;
int n,m,tot,sz,fa[N],size[N],q[N],pos[N],son[N],belong[N];
int deep[N],point[N],v[N],nxt[N],col[N];
struct data{
    int tr[N];
    int lowbit(int x){
        return x&(-x);
    }
    int query(int x){
        int ans=0;
        if (!x) return ans;
        for (int i=x;i>=1;i-=lowbit(i)) ans+=tr[i];
        return ans;
    }
    void change(int x,int v){
        if (!x) return;
        for (int i=x;i<=n;i+=lowbit(i)) tr[i]+=v; 
    }
}T[4];
void add(int x,int y)
{
    tot++; nxt[tot]=point[x]; point[x]=tot; v[tot]=y;
    tot++; nxt[tot]=point[y]; point[y]=tot; v[tot]=x;
}
void dfs(int x,int f)
{
    size[x]=1; deep[x]=deep[f]+1;
    for (int i=point[x];i;i=nxt[i]){
        if (v[i]==f) continue;
        fa[v[i]]=x;
        dfs(v[i],x);
        size[x]+=size[v[i]];
        if (size[son[x]]<size[v[i]]) son[x]=v[i];
    }
}
void dfs1(int x,int chain)
{
    pos[x]=++sz; belong[x]=chain; q[sz]=x;
    if (!son[x]) return;
    dfs1(son[x],chain);
    for (int i=point[x];i;i=nxt[i])
     if (v[i]!=fa[x]&&v[i]!=son[x])
      dfs1(v[i],v[i]);
}
int calc(int l,int r,int opt)
{
    int R=r; int ans=r;
    while (l<=r) {
        int mid=(l+r)/2;
        int t=T[3].query(R)-T[3].query(mid-1);
        int len=(opt==0?R-mid+1:0);
        if (t==len) r=mid-1,ans=min(ans,mid);
        else l=mid+1;
    }
    return ans;
}
int gettop(int x,int opt)
{
    int last=0;
    while (x) {
        int r=pos[x]; int l=pos[belong[x]];
        int t=T[3].query(r)-T[3].query(l-1);
        int len=(opt==0?r-l+1:0);
        if (t==len) {
         if (col[fa[belong[x]]]!=opt&&fa[belong[x]]) return belong[x];
         last=x,x=fa[belong[x]];
        }
        else {
            t=calc(l,r,opt);
            return q[t];
        }
    } 
    return 1;
}
void solve(int x,int y)
{
    if (!x||!y) return;
    while (belong[x]!=belong[y]) {
        if (deep[belong[x]]<deep[belong[y]]) swap(x,y);
        T[0].change(pos[belong[x]],1);
        T[0].change(pos[x]+1,-1);
        x=fa[belong[x]];
    }
    if(pos[x]>pos[y]) swap(x,y);
    T[0].change(pos[x],1);
    T[0].change(pos[y]+1,-1);
}
void toggle(int x,int y,int v1,int opt,int k)
{
    while (belong[x]!=belong[y]) {
        if (deep[belong[x]]<deep[belong[y]]) swap(x,y);
        T[opt].change(pos[belong[x]],k*v1);
        T[opt].change(pos[x]+1,-k*v1);
        x=fa[belong[x]];
    }
    if (pos[x]>pos[y]) swap(x,y);
    T[opt].change(pos[x],k*v1);
    T[opt].change(pos[y]+1,-k*v1);
}
int main()
{
    freopen("a.in","r",stdin);
    freopen("my.out","w",stdout);
    scanf("%d",&n);
    for (int i=1;i<n;i++) {
        int x,y; scanf("%d%d",&x,&y);
        add(x,y);
    }
    dfs(1,0); 
    dfs1(1,1);
    for (int i=1;i<=n;i++) T[3].change(pos[i],1);
    for (int i=1;i<=n;i++) solve(1,fa[i]),col[i]=0;;
    scanf("%d",&m);
    for (int i=1;i<=m;i++) {
        int opt,x; 
        scanf("%d%d",&opt,&x);
        if (opt==0) {
            int t=gettop(x,col[x]);
            printf("%d\n",T[col[x]].query(pos[t])+1);
        }
        else {
            int tt=gettop(x,col[x]);
            int v1=T[col[x]].query(pos[x])+1;
            int v2=T[col[x]^1].query(pos[x])+1;
            if (tt!=x) {
              toggle(tt==1?1:fa[tt],fa[x],v1,col[x],-1);
              T[col[x]^1].change(pos[fa[x]],v2);
              T[col[x]^1].change(pos[fa[x]]+1,-v2);
            }
            else 
            if (fa[x]){
                int t=gettop(fa[x],col[x]^1);
                toggle(t==1?1:fa[t],fa[x],v2,col[x]^1,1);
                T[col[x]].change(pos[fa[x]],-v1);
                T[col[x]].change(pos[fa[x]]+1,v1);
            }
            col[x]^=1; 
            if (col[x]) T[3].change(pos[x],-1);
            else T[3].change(pos[x],1);
        }
    }
}