bzoj3257 树的难题（树形DP）

Description

给出一个无根树。树有N个点，边有权值。每个点都有颜色，是黑色、白色、
灰色这三种颜色之一，称为一棵三色树。
可爱的 Alice觉得，一个三色树为均衡的，当且仅当，树中不含有黑色结点
或者含有至多一个白色节点。然而，给出的三色树可能并不满足这个性质。
所以，Alice打算删去若干条边使得形成的森林中每棵树都是均衡的，花费
的代价等于删去的边的权值之和。请你计算需要花费的代价最小是多少。
注意，输入文件包含多组测试数据。
Input

第一行包含一个正整数 T，表示有 T组测试数据。接下来依次是 T组测试数
据。每组测试数据的第一行包含一个正整数 N。
第二行包含 N个 0、1、2之一的整数，依次表示点 1到点 N的颜色。其中，
0 表示黑色，1表示白色，2表示灰色。
接下来 N-1行，每行为三个整数ui、vi、c i，表示一条权值等于ci的边(ui, vi)。
Output

输出 T行，每行一个整数，依次表示每组测试数据的答案。
Sample Input

1
5
0 1 1 1 0
1 2 5
1 3 3
5 2 5
2 4 16

Sample Output
10
HINT

花费 10的代价删去边(1, 2)和边(2, 5)。

对于 100%的数据：1 ≤ N ≤ 300 000 ，1 ≤ T ≤ 5 ，0 ≤ ci<=10^9

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

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

时间复杂度O(n)。
转载自：https://www.cnblogs.com/clrs97/p/5700005.html

#include
#define rep(i,n) for(int i=0;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\n",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;
}

自己写的：

#include
#include
#include
#include 
#define ll long long
#define inf 0x3f3f3f3f3f3f3f3f
#define maxn 300005
#define rep(i,a,b) for(register int i=a;i<=b;i++)
using namespace std;
struct Edge{
    int v, next;
    ll w;
}edge[maxn<<1];
int head[maxn], cnt, n, m, a[maxn];
ll f[maxn][2][3],h[2][3];

void up(ll &a,ll b){
    if(a>b) a = b;
}

int fix(int a,int b){
    return min(2,a+b);
}

void add_edge(int u, int v, ll w){
    edge[cnt].v = v;
    edge[cnt].w = w;
    edge[cnt].next = head[u];
    head[u] = cnt++;
}

void dfs(int x,int fa){
    int i, j;
    rep(A,0,1) rep(B,0,2) f[x][A][B] = inf;
    f[x][0][0] = 0;
    for(i = head[x];i!=-1;i=edge[i].next){
        int v = edge[i].v;
        if(v==fa) continue;
        dfs(v,x);
        rep(A,0,1) rep(B,0,2) h[A][B]=inf;
        rep(A,0,1) rep(B,0,2){
            if(f[x][A][B]!=inf){
                rep(C,0,1) rep(D,0,2){
                    if(f[v][C][D]!=inf){
                        up(h[A|C][fix(B,D)],f[x][A][B]+f[v][C][D]);
                        if(!C||D<2) up(h[A][B],f[x][A][B]+f[v][C][D]+edge[i].w);
                    }
                }
            }
        }
        rep(A,0,1) rep(B,0,2) f[x][A][B] = h[A][B];    
    }
    rep(A,0,1) rep(B,0,2) h[A][B] = inf;
    rep(A,0,1) rep(B,0,2) if(f[x][A][B]!=inf) up(h[A|!a[x]][fix(B,a[x]==1)],f[x][A][B]);
    rep(A,0,1) rep(B,0,2) f[x][A][B] = h[A][B];
}
int main(){
    int i, j, t, u, v;
    ll w;
    ll ans;
    for(scanf("%d",&t);t--;printf("%lld\n",ans)){
        scanf("%d",&n);
        ans = inf;
        cnt = 0;
        for(i = 1; i <= n; i++) {
            scanf("%d",&a[i]);
            head[i] = -1;
        }
        for(i = 1; i < n; i++) {
            scanf("%d%d%lld",&u,&v,&w);
            add_edge(u,v,w);
            add_edge(v,u,w);
        }
        dfs(1,0);
        rep(A,0,1) rep(B,0,2) if(!A||B<2) up(ans,h[A][B]);
    }
}