树的难题

题目大意

一颗树有n个结点，每个结点有黑白灰其中一种颜色。
要求你删去树中若干条边，使得形成的森林中每一颗树都没有白色结点或至多有一个黑色结点。删去一条边的代价为该边的边权，最小化代价。

树形DP

注意到一个合法的联通块的任意子联通块也符合条件。
所以树形DP，设f[i]，g[i]，h[i]表示将以i为根的子树分成合法森林，i结点所在联通块有0个黑点、1个黑点、0个白点的最小代价。
转移方程很容易了。具体见代码。
注意叶子节点如果不是黑色的g值为正无穷。
然后g的转移涉及需要从一个g转移过来的情况，可以先全部不从g转移过来，然后找到每个儿子更换成从g转移所造成影响，找到影响最大的那个，然后从其的g转移过来。
最后我爆栈了，懒得打人工。

#include<cstdio>
#include<algorithm>
#define fo(i,a,b) for(i=a;i<=b;i++)
using namespace std;
typedef long long ll;
const int maxn=300000+10;
const ll inf=1000000000000000;
ll f[maxn],g[maxn],h[maxn],dis[maxn*2];
int c[maxn],fi[maxn],go[maxn*2],next[maxn*2];
int i,j,k,l,t,n,m,tot,ca;
void add(int x,int y,int z){
    go[++tot]=y;
    dis[tot]=z;
    next[tot]=fi[x];
    fi[x]=tot;
}
void dfs(int x,int y){//0 black 1 white 2 gray
    bool leaf=1;
    f[x]=g[x]=h[x]=0;
    if (!c[x]) f[x]=inf;
    else if (c[x]==1) h[x]=inf;
    int t=fi[x],k=0,l=0;
    while (t){
        if (go[t]!=y){
            dfs(go[t],x);
            leaf=0;
            if (c[x]){
                f[x]+=min(f[go[t]],min(g[go[t]],h[go[t]])+dis[t]);
                if (!k||min(f[k],min(g[k],h[k])+l)-g[k]<min(f[go[t]],min(g[go[t]],h[go[t]])+dis[t])-g[go[t]]){
                    k=go[t];
                    l=dis[t];
                }
            }
            g[x]+=min(f[go[t]],min(g[go[t]],h[go[t]])+dis[t]);
            if (c[x]!=1) h[x]+=min(h[go[t]],min(f[go[t]],g[go[t]])+dis[t]);
        }
        t=next[t];
    }
    if (c[x]) g[x]=g[x]-min(f[k],min(g[k],h[k])+l)+g[k];
    if (leaf&&c[x]) g[x]=inf;
}
int read(){
    int x=0;
    char ch=getchar();
    while (ch<'0'||ch>'9') ch=getchar();
    while (ch>='0'&&ch<='9'){
        x=x*10+ch-'0';
        ch=getchar();
    }
    return x;
}
int main(){
    ca=read();
    while (ca--){
        n=read();
        fill(fi+1,fi+n+1,0);
        tot=0;
        fo(i,1,n) c[i]=read();
        fo(i,1,n) 
            if (c[i]<2) c[i]=1-c[i];
        fo(i,1,n-1){
            j=read();k=read();l=read();
            add(j,k,l);add(k,j,l);
        }
        dfs(1,0);
        printf("%lld\n",min(f[1],min(g[1],h[1])));
    }
}