BZOJ 3257: 树的难题

树形DP

#include<cstdio>

#include<algorithm>

#define rep(i,x,y) for (int i=x; i<=y; i++)

using namespace std;

int cnt,last[300005],c[300005];

long long F[300005][2][3],G[2][3];

struct node{

	int to,next,val;

}e[600005];

void add(int a,int b,int c){

	e[++cnt].to=b;

	e[cnt].next=last[a];

	e[cnt].val=c;

	last[a]=cnt;

}

void dfs(int x,int fa){

	rep(i,0,1) rep(j,0,2) F[x][i][j]=1ll<<60;

	F[x][c[x]==0][c[x]==1]=0;

	for (int i=last[x]; i; i=e[i].next){

		int V=e[i].to;

		if (V==fa) continue;

		dfs(V,x);

		rep(nowx,0,1) rep(nowy,0,2) G[nowx][nowy]=1ll<<60;

		rep(prex,0,1) rep(prey,0,2)

		rep(nowx,0,1) rep(nowy,0,2){

			int tox=min(1,prex+nowx),toy=min(2,prey+nowy);

			if (tox!=1 || toy!=2) G[tox][toy]=min(G[tox][toy],F[x][prex][prey]+F[V][nowx][nowy]);

			G[prex][prey]=min(G[prex][prey],F[x][prex][prey]+F[V][nowx][nowy]+e[i].val);

		}

		rep(nowx,0,1) rep(nowy,0,2) F[x][nowx][nowy]=G[nowx][nowy];

	}

}

int main(){

	int T;

	scanf("%d",&T);

	while (T--){

		int n;

		scanf("%d",&n);

		for (int i=1; i<=n; i++) scanf("%d",&c[i]);

		cnt=0;

		for (int i=1; i<=n; i++) last[i]=0;

		for (int i=1; i<n; i++){

			int x,y,z;

			scanf("%d%d%d",&x,&y,&z);

			add(x,y,z);

			add(y,x,z);

		}

		dfs(1,0);

		long long ans=1ll<<60;

		rep(i,0,1) rep(j,0,2) ans=min(ans,F[1][i][j]);

		printf("%lld\n",ans);

	}

	return 0;

}