題目

題目連結：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;
}