BZOJ3257 : 树的难题

设$f[x][i][j]$表示以$x$为根的子树，与$x$连通部分有$i$个黑点，$j$个白点，不联通部分都是均衡的最小代价。若$i>1$，则视作$1$；若$j>2$，则视作$2$。

然后进行树形DP即可，转移的时候如果不要那棵子树，那么那棵子树的状态必须满足$!i||j<2$。

时间复杂度$O(n)$。

#include<cstdio>
#define rep(i,n) for(int i=0;i<n;i++)
typedef long long ll;
const int N=300010;
const ll inf=1LL<<60;
int T,n,i,x,y,z,a[N],g[N],v[N<<1],w[N<<1],nxt[N<<1],ed;
ll f[N][2][3],h[2][3],ans;
inline void read(int&a){char c;while(!(((c=getchar())>='0')&&(c<='9')));a=c-'0';while(((c=getchar())>='0')&&(c<='9'))(a*=10)+=c-'0';}
inline void up(ll&a,ll b){if(a>b)a=b;}
inline void add(int x,int y,int z){v[++ed]=y;w[ed]=z;nxt[ed]=g[x];g[x]=ed;}
inline int fix(int x){return x<2?x:2;}
void dfs(int x,int y){
  rep(A,2)rep(B,3)f[x][A][B]=inf;
  f[x][0][0]=0;
  for(int i=g[x];i;i=nxt[i])if(v[i]!=y){
    int u=v[i];
    dfs(u,x);
    rep(A,2)rep(B,3)h[A][B]=inf;
    rep(A,2)rep(B,3)if(f[x][A][B]<inf)rep(C,2)rep(D,3)if(f[u][C][D]<inf){
      up(h[A|C][fix(B+D)],f[x][A][B]+f[u][C][D]);
      if(!C||D<2)up(h[A][B],f[x][A][B]+f[u][C][D]+w[i]);
    }
    rep(A,2)rep(B,3)f[x][A][B]=h[A][B];
  }
  rep(A,2)rep(B,3)h[A][B]=inf;
  rep(A,2)rep(B,3)if(f[x][A][B]<inf)up(h[A|!a[x]][fix(B+(a[x]==1))],f[x][A][B]);
  rep(A,2)rep(B,3)f[x][A][B]=h[A][B];
}
int main(){
  for(read(T);T--;printf("%lld
",ans)){
    read(n);
    for(ed=0,i=1;i<=n;i++)read(a[i]),g[i]=0;
    for(i=1;i<n;i++)read(x),read(y),read(z),add(x,y,z),add(y,x,z);
    dfs(1,0);
    ans=inf;
    rep(A,2)rep(B,3)if(!A||B<2)up(ans,h[A][B]);
  }
  return 0;
}