3257: 树的难题

3257: 树的难题

https://www.lydsy.com/JudgeOnline/problem.php?id=3257

分析：

　　状态只与黑点有0个，大于0个，和白点有0个，1个，大于1个这六个状态有关系。f[u][0/1][0/1/2]表示以u为根的子树最小花费。

　　转移方程有点难写！！！

代码：

#include<cstdio>
#include<algorithm>
#include<cstring>
#include<cmath>
#include<iostream>
#include<cctype>
#include<set>
#include<vector>
#include<queue>
#include<map>
using namespace std;
typedef long long LL;

inline int read() {
    int x=0,f=1;char ch=getchar();for(;!isdigit(ch);ch=getchar())if(ch=='-')f=-1;
    for(;isdigit(ch);ch=getchar())x=x*10+ch-'0';return x*f;
}

const int N = 300100;
const LL INF = 1e18;

int head[N], nxt[N << 1], to[N << 1], len[N << 1], col[N];
LL f[N][2][3], g[2][3];
int En;

void add_edge(int u,int v,int w) {
    ++En; to[En] = v; len[En] = w; nxt[En] = head[u]; head[u] = En;
    ++En; to[En] = u; len[En] = w; nxt[En] = head[v]; head[v] = En;
}

void update(int u,int v,int w) {
    for (int a=0; a<=1; ++a) 
        for (int b=0; b<=2; ++b) g[a][b] = INF;
    for (int a=0; a<=1; ++a) // 原来树上的黑点 
        for (int b=0; b<=2; ++b) // 原来书上的白点 
            if (f[u][a][b] != INF) 
                for (int c=0; c<=1; ++c) // 子树的黑点 
                    for (int d=0; d<=2; ++d) // 子树的白点 
                        if (f[v][c][d] != INF) {
                            int p1 = a + c >= 1 ? 1 : 0, p2 = b + d >= 2 ? 2 : b + d;  
                            g[p1][p2] = min(g[p1][p2], f[u][a][b] + f[v][c][d]); // 原来树上的+子树上的 
                            if (!c || d < 2) g[a][b] = min(g[a][b], f[u][a][b] + f[v][c][d] + w); // 断开这条边，使得子树合法 
                        }
    memcpy(f[u], g, sizeof(g));
}

void dfs(int u,int fa) {
    for (int a=0; a<=1; ++a) 
        for (int b=0; b<=2; ++b) f[u][a][b] = INF;
    f[u][col[u]==0][col[u]==1] = 0;
    for (int i=head[u]; i; i=nxt[i]) {
        int v = to[i];
        if (v == fa) continue;
        dfs(v, u);
        update(u, v, len[i]);
    }
}

void solve() {
    int n = read();
    En = 0;
    memset(head, 0, sizeof(head)); //---
    for (int i=1; i<=n; ++i) col[i] = read();
    for (int i=1; i<n; ++i) {
        int u = read(), v = read(), w = read();
        add_edge(u, v, w);
    }
    dfs(1, 0);
    LL Ans = min(min(f[1][0][0], f[1][0][1]), f[1][0][2]);
    Ans = min(Ans, min(f[1][1][0], f[1][1][1]));
    cout << Ans << "\n";
}

int main() {
    int T = read();
    while (T --) solve();
    return 0;
}