【JZOJ3347】树的难题【树形dp】【拓扑排序】

题目大意：

题目链接：https://jzoj.net/senior/#main/show/3347

---

思路：

设$f[x]$表示以$x$为根的子树中，满足有0个黑色点，若干个白色点的最少切割次数；$g[x]$表示以$x$为根的子树中，满足有若干个黑色点，0个白色点的最少切割次数；$h[x]$表示以$x$为根的子树中，满足有若干个黑色点，1个白色点的最少切割次数。$i$为$x$的儿子。

• 如果$x$是黑色的
  1. $f[x]$需要满足没有黑色节点，显然不满足。
  2. $g[x]={\sum }^{i}min(f[v]+e[i].dis,g[v],h[v]+e[i].dis)$
  3. $h[x]=min(h[u],h[v]+g[u]-minn(f[v]+e[i].dis,g[v],h[v]+e[i].dis))$
• 如果$x$是白色的
  1. $f[x]={\sum }^{i}min(f[v],g[v]+e[i].dis,h[v]+e[i].dis)$
  2. $g[x]$显然不满足
  3. $h[x]={\sum }^{i}min(f[v]+e[i].dis,g[v],h[v]+e[i].dis)$
• 如果$x$是灰色的
  1. $f[x]={\sum }^{i}min(f[v],g[v]+e[i].dis,h[v]+e[i].dis)$
  2. $g[x]={\sum }^{i}min(f[v]+e[i].dis,g[v],h[v]+e[i].dis)$
  3. $h[x]=min(h[u],h[v]+g[u]-minn(f[v]+e[i].dis,g[v],h[v]+e[i].dis))$

还是比较好推的。毕竟除了设状态以外都是自己推的。
代码又长又丑，恶心死了。听说普通的$dp$会爆栈，还敲了一个$topsort$来$bfs$做$dp$，而且$topsort$和$bfs$还是一起进行的。。。
时间复杂度$O(Tn)$

---

代码：

#include <queue>
#include <cstdio>
#include <string>
#include <cstring>
using namespace std;
typedef long long ll;

const int N=300010;
const ll Inf=1e17;
int T,n,tot,root,head[N],col[N],in[N];
ll f[N],g[N],h[N];

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

int read()
{
	int d=0;
	char ch=getchar();
	while (!isdigit(ch)) ch=getchar();
	while (isdigit(ch))
		d=(d<<3)+(d<<1)+ch-48,ch=getchar();
	return d;
}

ll minn(ll x,ll y,ll z)
{
	return min(x,min(y,z));
}

void add(int from,int to,ll dis)
{
	e[++tot].to=to;
	e[tot].dis=dis;
	e[tot].next=head[from];
	head[from]=tot;
}

void dp()
{
	queue<int> q;
	for (int i=1;i<=n;i++)
		if (in[i]==1) q.push(i);
	while (q.size())
	{
		int u=q.front();
		q.pop();
		root=u;
		in[u]=0;
		
		if (col[u]==0)
		{
			f[u]=Inf;
			for (int i=head[u];~i;i=e[i].next)
			{
				int v=e[i].to;
				if (!in[v]) g[u]+=minn(f[v]+e[i].dis,g[v],h[v]+e[i].dis);
			}
			for (int i=head[u];~i;i=e[i].next)
			{
				int v=e[i].to;
				if (!in[v]) h[u]=min(h[u],h[v]+g[u]-minn(f[v]+e[i].dis,g[v],h[v]+e[i].dis));
			}
		}
		
		if (col[u]==1)
		{
			g[u]=Inf;
			for (int i=head[u];~i;i=e[i].next)
			{
				int v=e[i].to;
				if (!in[v])
				{
					f[u]+=minn(f[v],g[v]+e[i].dis,h[v]+e[i].dis);
					h[u]+=minn(f[v]+e[i].dis,g[v],h[v]+e[i].dis);
				}
			}
		}
		
		if (col[u]==2)
		{
			for (int i=head[u];~i;i=e[i].next)
			{
				int v=e[i].to;
				if (!in[v])
				{
					f[u]+=minn(f[v],g[v]+e[i].dis,h[v]+e[i].dis);
					g[u]+=minn(f[v]+e[i].dis,g[v],h[v]+e[i].dis);
				}
			}
			for (int i=head[u];~i;i=e[i].next)
			{
				int v=e[i].to;
				if (!in[v]) h[u]=min(h[u],h[v]+g[u]-minn(f[v]+e[i].dis,g[v],h[v]+e[i].dis));
			}
		}
		
		for (int i=head[u];~i;i=e[i].next)
		{
			int v=e[i].to;
			if (in[v]>1)
			{
				in[v]--;
				if (in[v]==1) q.push(v);
			}
		}
	}
}

int main()
{
	T=read();
	while (T--)
	{
		memset(head,-1,sizeof(head));
		memset(in,0,sizeof(in));
		tot=0;
		n=read();
		for (int i=1;i<=n;i++)
		{
			col[i]=read();
			if (col[i]==0) f[i]=Inf,g[i]=0,h[i]=Inf;
			if (col[i]==1) f[i]=0,g[i]=Inf,h[i]=0;
			if (col[i]==2) f[i]=0,g[i]=0,h[i]=Inf;
		}
		for (int i=1,x,y,z;i<n;i++)
		{
			x=read(); y=read(); z=read();
			add(x,y,(ll)z); add(y,x,(ll)z);
			in[x]++; in[y]++;
		}
		dp();
		printf("%lld\n",minn(f[root],g[root],h[root]));
	}
	return 0;
}