【CF343D】Water Tree（树链剖分）

大水题

#include<bits/stdc++.h>
const int N=500005;
using namespace std;
template<class T>
inline void read(T &x)
{
	x=0; int f=1;
	static char ch=getchar();
	while(!isdigit(ch)) {if(ch=='-') f=-1;ch=getchar();}
	while(isdigit(ch)) x=x*10+ch-'0',ch=getchar();
	x*=f;
}
int n,tot,first[N],Q;
struct Edge
{
	int to,next;
}edge[2*N];
inline void addedge(int u,int v)
{
	tot++;
	edge[tot].to=v; edge[tot].next=first[u]; first[u]=tot;
}
int father[N],size[N],maxson[N];
void dfs1(int now,int fa)
{
	father[now]=fa;
	size[now]=1;
	for(int u=first[now];u;u=edge[u].next)
	{
		int vis=edge[u].to;
		if(vis==fa)	continue;
		dfs1(vis,now);
		size[now]+=size[vis];
		if(size[vis]>size[maxson[now]])	maxson[now]=vis;
	}
}
int top[N],id[N],cnt;
void dfs2(int now,int topf)
{
	top[now]=topf;
	id[now]=++cnt;
	if(maxson[now])	dfs2(maxson[now],topf);
	for(int u=first[now];u;u=edge[u].next)
	{
		int vis=edge[u].to;
		if(vis==father[now]||vis==maxson[now])	continue;
		dfs2(vis,vis);
	}
}
struct SegMentTree
{
	int l,r,sum,f;
}tree[4*N];
inline void build(int now,int l,int r)
{
	tree[now].l=l; tree[now].r=r; tree[now].f=-1;
	if(l==r)	return;
	int mid=(l+r)>>1;
	build(now<<1,l,mid);
	build(now<<1|1,mid+1,r);
}
inline void pushdown(int now)
{
	if(tree[now].f!=-1)
	{
		int f=tree[now].f;
		tree[2*now].f=tree[2*now+1].f=f;
		tree[2*now].sum=(tree[2*now].r-tree[2*now].l+1)*f;
		tree[2*now+1].sum=(tree[2*now+1].r-tree[2*now+1].r+1)*f;
		tree[now].f=-1;
	}
}
inline void Cover(int now,int l,int r,int c)
{
	if(l<=tree[now].l&&tree[now].r<=r) 
	{
		tree[now].f=c,tree[now].sum=(tree[now].r-tree[now].l+1)*c;
		return;
	}
	pushdown(now);
	int mid=(tree[now].l+tree[now].r)>>1;
	if(l<=mid)	Cover(2*now,l,r,c);
	if(r>mid) Cover(2*now+1,l,r,c);
}
inline int Query(int now,int p)
{
	if(tree[now].l==tree[now].r)	return tree[now].sum;
	pushdown(now);
	int mid=(tree[now].l+tree[now].r)>>1,ans;
	if(p<=mid)	ans=Query(2*now,p);
	else ans=Query(2*now+1,p);
	return ans;
}
inline void Cover_Sub_Tree(int u){Cover(1,id[u],id[u]+size[u]-1,1);} 
inline void Cover_To_Root(int u)
{
	while(top[u]!=1)
	{
		Cover(1,id[top[u]],id[u],0);
		u=father[top[u]];
	}
	Cover(1,1,id[u],0);
}
int main()
{
	read(n);
	for(int i=1,u,v;i<=n-1;i++) read(u),read(v),addedge(u,v),addedge(v,u);
	dfs1(1,0); dfs2(1,1); build(1,1,n);
	read(Q);
	int opt,u;
	while(Q--)
	{
		read(opt),read(u);
		if(opt==1) Cover_Sub_Tree(u);
		if(opt==2) Cover_To_Root(u);
		if(opt==3) cout<<Query(1,id[u])<<'\n';
	}
	return 0;
}