【CF1626E】Black and White Tree

题目

题目链接：https://codeforces.com/contest/1626/problem/E
一棵 \(n\) 个节点的树，节点有黑白两种颜色，黑点的数量至少为 \(2\)。
一开始有一个石子在树的某一个点上，每一次操作你可以选择一个黑点，让石子向黑点的方向走一条边。不能连续两次操作选择相同的黑点。
对于每一个点，求出如果石子在这个点上，是否有方法可以将石子移动到黑点上。
\(n\leq 3\times 10^5\)。

思路

设点 \(1\) 为根，考虑石子在点 \(1\) 时怎样才能被移动到黑点上。

• 如果点 \(1\) 周围有黑点，或者点 \(1\) 就是黑点，那么显然是有解的。
• 如果存在两个黑点成祖孙关系，那么只需要不断选择这两个黑点即可。
• 如果不存在两个黑点成祖孙关系，那么如果石子要到达黑点 \(x\)，最后一步只能选择 \(x\) 来移动，那么就必须通过其他黑点到达 \(x\) 的父亲，也就是说 \(x\) 的父亲的子树内必须有其他黑点。
  而若 \(x\) 的父亲的子树内有其他黑点 \(y\)，那么一定可以移动到 \(x\)，只需要不断选择 \(x\) 和 \(y\)，到达 \(x\) 的父亲后再选择一次 \(x\) 即可。

后面两点等价于把每一个黑点的父亲拿出来，然后存在两个父亲他们互为祖孙关系。
也就是把每个黑点父亲的子树打上 \(+1\) 的标记后，存在一个节点标记至少为 \(2\)。
直接线段树维护全局最大值，然后换根的时候维护一下即可。
时间复杂度 \(O(n\log n)\)。

代码

#include <bits/stdc++.h>
using namespace std;

const int N=300010;
int n,tot,head[N],a[N],b[N],id[N],siz[N],L[N],R[N];
bool ans[N];

struct edge
{
	int next,to;
}e[N*2];

void add(int from,int to)
{
	e[++tot]=(edge){head[from],to};
	head[from]=tot;
}

struct SegTree
{
	int mx[N*4],lazy[N*4];
	
	void pushdown(int x)
	{
		if (!lazy[x]) return;
		mx[x*2]+=lazy[x]; lazy[x*2]+=lazy[x];
		mx[x*2+1]+=lazy[x]; lazy[x*2+1]+=lazy[x];
		lazy[x]=0;
	}
	
	void update(int x,int l,int r,int ql,int qr,int v)
	{
		if (ql>qr) return;
		if (ql<=l && qr>=r) { mx[x]+=v; lazy[x]+=v; return; }
		pushdown(x);
		int mid=(l+r)>>1;
		if (ql<=mid) update(x*2,l,mid,ql,qr,v);
		if (qr>mid) update(x*2+1,mid+1,r,ql,qr,v);
		mx[x]=max(mx[x*2],mx[x*2+1]);
	}
}seg;

void dfs1(int x,int fa)
{
	id[x]=++tot; siz[x]=1;
	for (int i=head[x];~i;i=e[i].next)
	{
		int v=e[i].to;
		b[x]+=a[v];
		if (v!=fa) dfs1(v,x),siz[x]+=siz[v];
	}
	L[x]=id[x]; R[x]=id[x]+siz[x]-1;
	if (x!=1) seg.update(1,1,n,L[x],R[x],b[x]-a[fa]);
}

void dfs2(int x,int fa)
{
	ans[x]=a[x]|b[x]|(seg.mx[1]>=2);
	for (int i=head[x];~i;i=e[i].next)
	{
		int v=e[i].to;
		if (v!=fa)
		{
			seg.update(1,1,n,L[v],R[v],-(b[v]-a[x]));
			seg.update(1,1,n,1,L[v]-1,b[x]-a[v]);
			seg.update(1,1,n,R[v]+1,n,b[x]-a[v]);
			dfs2(v,x);
			seg.update(1,1,n,L[v],R[v],b[v]-a[x]);
			seg.update(1,1,n,1,L[v]-1,-(b[x]-a[v]));
			seg.update(1,1,n,R[v]+1,n,-(b[x]-a[v]));
		}
	}
}

int main()
{
	memset(head,-1,sizeof(head));
	scanf("%d",&n);
	for (int i=1;i<=n;i++) scanf("%d",&a[i]);
	for (int i=1,x,y;i<n;i++)
	{
		scanf("%d%d",&x,&y);
		add(x,y); add(y,x);
	}
	tot=0;
	dfs1(1,0); dfs2(1,0);
	for (int i=1;i<=n;i++)
		printf("%d ",ans[i]);
	return 0;
}