3257: 树的难题

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

分析：

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

　　转移方程有点难写！！！

代码：

 1 #include<cstdio>
 2 #include<algorithm>
 3 #include<cstring>
 4 #include<cmath>
 5 #include<iostream>
 6 #include<cctype>
 7 #include<set>
 8 #include<vector>
 9 #include<queue>
10 #include<map>
11 using namespace std;
12 typedef long long LL;
13 
14 inline int read() {
15     int x=0,f=1;char ch=getchar();for(;!isdigit(ch);ch=getchar())if(ch=='-')f=-1;
16     for(;isdigit(ch);ch=getchar())x=x*10+ch-'0';return x*f;
17 }
18 
19 const int N = 300100;
20 const LL INF = 1e18;
21 
22 int head[N], nxt[N << 1], to[N << 1], len[N << 1], col[N];
23 LL f[N][2][3], g[2][3];
24 int En;
25 
26 void add_edge(int u,int v,int w) {
27     ++En; to[En] = v; len[En] = w; nxt[En] = head[u]; head[u] = En;
28     ++En; to[En] = u; len[En] = w; nxt[En] = head[v]; head[v] = En;
29 }
30 
31 void update(int u,int v,int w) {
32     for (int a=0; a<=1; ++a) 
33         for (int b=0; b<=2; ++b) g[a][b] = INF;
34     for (int a=0; a<=1; ++a) // 原来树上的黑点 
35         for (int b=0; b<=2; ++b) // 原来书上的白点 
36             if (f[u][a][b] != INF) 
37                 for (int c=0; c<=1; ++c) // 子树的黑点 
38                     for (int d=0; d<=2; ++d) // 子树的白点 
39                         if (f[v][c][d] != INF) {
40                             int p1 = a + c >= 1 ? 1 : 0, p2 = b + d >= 2 ? 2 : b + d;  
41                             g[p1][p2] = min(g[p1][p2], f[u][a][b] + f[v][c][d]); // 原来树上的+子树上的 
42                             if (!c || d < 2) g[a][b] = min(g[a][b], f[u][a][b] + f[v][c][d] + w); // 断开这条边，使得子树合法 
43                         }
44     memcpy(f[u], g, sizeof(g));
45 }
46 
47 void dfs(int u,int fa) {
48     for (int a=0; a<=1; ++a) 
49         for (int b=0; b<=2; ++b) f[u][a][b] = INF;
50     f[u][col[u]==0][col[u]==1] = 0;
51     for (int i=head[u]; i; i=nxt[i]) {
52         int v = to[i];
53         if (v == fa) continue;
54         dfs(v, u);
55         update(u, v, len[i]);
56     }
57 }
58 
59 void solve() {
60     int n = read();
61     En = 0;
62     memset(head, 0, sizeof(head)); //---
63     for (int i=1; i<=n; ++i) col[i] = read();
64     for (int i=1; i<n; ++i) {
65         int u = read(), v = read(), w = read();
66         add_edge(u, v, w);
67     }
68     dfs(1, 0);
69     LL Ans = min(min(f[1][0][0], f[1][0][1]), f[1][0][2]);
70     Ans = min(Ans, min(f[1][1][0], f[1][1][1]));
71     cout << Ans << "
";
72 }
73 
74 int main() {
75     int T = read();
76     while (T --) solve();
77     return 0;
78 }