JZOJ 3347. 【NOI2013模拟】树的难题【树形dp】

---

题目：

传送门

---

题意：

给出一棵含有三色点的森林，让我们删去一部分边，使得剩下的每棵树都是符合条件而均衡的

---

分析：

直接上树形$dp$
用$f_{i,0—5}$表示以$i$为根的子树的五种状态

---

代码：

#include<cstdio>
#include<string>
#include<cstring>
#include<iostream>
#include<algorithm>
#include<queue>
#define LL long long 
#define min3(a1,a2,a3) min(min(a1,a2),a3)
#define min4(a1,a2,a3,a4) min(min(a1,a2),min(a3,a4))
#define min5(a1,a2,a3,a4,a5) min(min(min(a1,a2),min(a3,a4)),a5)
using namespace std;
const LL LZX=1e18;
inline LL read() {
    LL d=0,f=1;char s=getchar();
    while(s<'0'||s>'9'){if(s=='-')f=-1;s=getchar();}
    while(s>='0'&&s<='9'){d=d*10+s-'0';s=getchar();}
    return d*f;
}  
using namespace std;
struct node{
	LL to,next,w;
}a[600005];
LL t,n,c[300005],ls[300005],tot,f[300005][5],q[300005],fa[300005];
void add(LL x,LL y,LL w)
{
	a[++tot]=(node){y,ls[x],w};
	ls[x]=tot;
	return;
}
void dfs()
{
	LL head=1,tail=0;
	q[++tail]=1;
	while(head<=tail)
	{
		LL u=q[head++];
		for(LL i=ls[u];i;i=a[i].next)
		{
			LL v=a[i].to;
			if(v==fa[u]) continue;
			fa[v]=u;q[++tail]=v;
		}
	}
	for(LL i=tail;i>=1;i--)
	{
		LL u=q[i];
		LL white1=0,white2=0;
		for(LL i=ls[u];i;i=a[i].next)
		{
			LL v=a[i].to,z;
			if(v==fa[u]) continue;
			f[u][0]+=min(min4(f[v][1],f[v][2],f[v][3],f[v][4])+a[i].w,f[v][0]);
			f[u][1]+=min(min3(f[v][0],f[v][1],f[v][2]),min(f[v][3],f[v][4])+a[i].w);
			f[u][2]+=(z=min(f[v][0],min4(f[v][1],f[v][2],f[v][3],f[v][4])+a[i].w));
			white1=max(white1,z-f[v][2]);
			f[u][3]+=(z=min(min(f[v][0],f[v][4]),min3(f[v][3],f[v][2],f[v][1])+a[i].w));
			white2=max(white2,z-min(f[v][3],f[v][2]));
			f[u][4]+=min(min(f[v][0],f[v][4]),min3(f[v][1],f[v][2],f[v][3])+a[i].w);
		}
		f[u][2]-=white1;f[u][3]-=white2;
		if(c[u]==1)
		{
			f[u][3]=f[u][4];f[u][2]=f[u][0];
			f[u][1]=min3(f[u][0],f[u][1],f[u][2]);
			f[u][0]=f[u][4]=LZX;
		}
		else if(c[u]==0)
		{
			f[u][3]=min3(f[u][3],f[u][2],f[u][0]);
			f[u][0]=f[u][1]=f[u][2]=LZX;
		}
		f[u][1]=min3(f[u][1],f[u][0],f[u][2]);
		f[u][2]=min(f[u][2],f[u][0]);
		f[u][3]=min3(f[u][3],f[u][2],f[u][0]);
		f[u][4]=min(f[u][4],f[u][0]);
	}
	return;
}
int main()
{
	t=read();
	while(t--)
	{
		memset(f,0,sizeof(f));memset(ls,0,sizeof(ls));tot=0;
		n=read();
		for(LL i=1;i<=n;i++) c[i]=read();
		for(LL i=1;i<n;i++)
		{
		  LL x=read(),y=read(),w=read();
		  add(x,y,w);add(y,x,w);
		}
		dfs();
		printf("%lld\n",min5(f[1][0],f[1][1],f[1][2],f[1][3],f[1][4]));
	}
	return 0;
}