Codeforces Round #200 (Div. 1)D. Water Tree

简单的树链剖分+线段树

#include<bits\stdc++.h>
using namespace std;
#define pb push_back
#define lson root<<1,l,midd
#define rson root<<1|1,midd+1,r
const int M=5e5+5;
vector<int>g[M];
int tree[M<<2],lazy[M<<2],top[M],son[M],fa[M],sz[M],dfn[M],to[M],deep[M],cnt,n;
void dfs1(int u,int f){
    sz[u]=1;
    fa[u]=f;
    deep[u]=deep[f]+1;
    for(int i=0;i<g[u].size();i++){
        int v=g[u][i];
        if(v!=f){
            dfs1(v,u);
            sz[u]+=sz[v];
            if(sz[v]>sz[son[u]])
                son[u]=v;
        }
    }
}
void dfs2(int u,int t){
    dfn[u]=++cnt;
    to[cnt]=u;
    top[u]=t;
    if(!son[u])
        return ;
    dfs2(son[u],t);
    for(int i=0;i<g[u].size();i++){
        int v=g[u][i];
        if(v!=fa[u]&&v!=son[u])
            dfs2(v,v);
    }
}
void pushdown(int root){
    tree[root<<1]=tree[root<<1|1]=lazy[root]-1;
    lazy[root<<1]=lazy[root<<1|1]=lazy[root];
    lazy[root]=0;
    
}
void update(int L,int R,int w,int root,int l,int r){
    if(L<=l&&r<=R){
        tree[root]=w;
        lazy[root]=w+1;
        return ;
    }
    if(lazy[root])
        pushdown(root);
    int midd=(l+r)>>1;
    if(L<=midd)
        update(L,R,w,lson);
    if(R>midd)
        update(L,R,w,rson);
}
void Update(int u,int v,int w){
    while(top[u]!=top[v]){
        if(deep[top[u]]<deep[top[v]])
            swap(u,v);
        update(dfn[top[u]],dfn[u],w,1,1,n);
        u=fa[top[u]];
    }
    if(deep[u]<deep[v])
        swap(u,v);
    update(dfn[v],dfn[u],w,1,1,n);
} 
int query(int pos,int root,int l,int r){
    if(l==r){
        return tree[root];
    }
    if(lazy[root])
        pushdown(root);
    int midd=(l+r)>>1;
    if(pos<=midd)
        return query(pos,lson);
    if(pos>midd)
        return query(pos,rson);
}
int main(){
    scanf("%d",&n);
    for(int i=1;i<n;i++){
        int u,v;
        scanf("%d%d",&u,&v);
        g[u].pb(v);
        g[v].pb(u);
    }
    dfs1(1,1);
    dfs2(1,1);
    int m;
    scanf("%d",&m);
    while(m--){
        int op,u;
        scanf("%d%d",&op,&u);
        if(op==1){
            update(dfn[u],dfn[u]+sz[u]-1,1,1,1,n);
        }
        else if(op==2){
            Update(1,u,0);
        }
        else{
            printf("%d\n",query(dfn[u],1,1,n));
        }
    }
    return 0;
}