CodeForces 343D : Water Tree 树链剖分

传送门

题目描述

给定 n 个点的一棵有根树，1 号点是根，维护三种操作： 1.把结点 v 及其子树赋值 1。 2.把结点 v 及其到根上的路径赋值 0。 3.询问 v 点的值。 初始时所有的点都是 0。

分析

学过树链剖分之后我们可以很快的秒掉这道题

代码

#include <bits/stdc++.h>

using namespace std;

const int N = 100010;
int h[N],ne[N],e[N],idx;
int cnt,id[N];
int dep[N],son[N],fa[N],top[N],sz[N];
int n,m;

struct Node{
    int l,r;
    int sum,add;
}tr[N  << 2];

void add(int x,int y){
    ne[idx] = h[x],e[idx] = y,h[x] = idx++; 
}

void dfs(int u,int f,int d){
    dep[u] = d,fa[u] = f,sz[u] = 1;
    for(int i = h[u];~i;i = ne[i]){
        int j = e[i];
        if(j == f) continue;
        dfs(j,u,d + 1);
        sz[u] += sz[j];
        if(sz[j] > sz[son[u]]) son[u] = j;
    }
}

void dfs(int u,int t){
    id[u] = ++cnt,top[u] = t;
    if(!son[u]) return;
    dfs(son[u],t);
    for(int i = h[u];~i;i = ne[i]){
        int j = e[i];
        if(j == fa[u] || j == son[u]) continue;
        dfs(j,j);
    }
}
void pushup(int u){
    tr[u].sum = tr[u << 1].sum + tr[u << 1 | 1].sum;
}

void pushdown(int u){
    auto &root = tr[u], &left = tr[u << 1], &right = tr[u << 1 | 1];
    if(root.add != -1){
        left.sum = root.add * (left.r - left.l + 1);
        left.add = root.add;
        right.sum = root.add * (right.r - right.l + 1);
        right.add = root.add;
        root.add = -1;
    }
}

void build(int u,int l,int r){
    tr[u] = {l,r,0,-1};
    if(l == r) return;
    int mid = l + r >> 1;
    build(u << 1,l,mid),build(u << 1 | 1,mid + 1,r);
}

void modify(int u,int l,int r,int k){
    if(tr[u].l >= l && tr[u].r <= r){
        tr[u].add = k;
        tr[u].sum = k * (tr[u].r - tr[u].l + 1);
        return;
    }
    pushdown(u);
    int mid = tr[u].l + tr[u].r >> 1;
    if(l <= mid) modify(u << 1,l,r,k);
    if(r > mid) modify(u << 1 | 1,l,r,k);
    pushup(u);
}

void modify_path(int u,int v,int k){
    while(top[u] != top[v]){
        if(dep[top[u]] < dep[top[v]]) swap(u,v);
        modify(1,id[top[u]],id[u],k);
        u = fa[top[u]];
    }
    if(dep[u] < dep[v]) swap(u,v);
    modify(1,id[v],id[u],k);
}

void modify_tree(int u,int k){
    modify(1,id[u],id[u] + sz[u] - 1,k);
}

int query(int u,int l,int r){
    if(tr[u].l >= l && tr[u].r <= r){
        return tr[u].sum;
    }
    pushdown(u);
    int mid = tr[u].l + tr[u].r >> 1;
    int sum = 0;
    if(l <= mid) sum = query(u << 1,l,r);
    if(r > mid) sum += query(u << 1 | 1,l,r);
    return sum;
}

int main(){
    scanf("%d",&n);
    memset(h,-1,sizeof h);
    for(int i = 2;i <= n;i++){
        int x,y;
        scanf("%d%d",&x,&y);
        add(x,y),add(y,x);
    }
    dfs(1,-1,1);
    dfs(1,1);
    build(1,1,n);
    scanf("%d",&m);
    int x,op;
    while(m--){
        scanf("%d%d",&op,&x);
        if(op == 1) modify_tree(x,1);
        if(op == 2) modify_path(1,x,0);
        if(op == 3) printf("%d\n",query(1,id[x],id[x]));
    }

}