diff --git a/TangDou/Topic/HuanGenDp/CF1324F.cpp b/TangDou/Topic/HuanGenDp/CF1324F.cpp index db72c1d..5fa8a36 100644 --- a/TangDou/Topic/HuanGenDp/CF1324F.cpp +++ b/TangDou/Topic/HuanGenDp/CF1324F.cpp @@ -1,60 +1,54 @@ -#include -#include -#include -#include +#include using namespace std; -const int N = 2e5 + 10; +const int N = 2e5 + 10, M = N << 1; -struct Edge { - int to, next; -} edge[N << 1]; -int idx; -int h[N]; - -void add_edge(int u, int v) { - edge[++idx] = {v, h[u]}; - h[u] = idx; +// 链式前向星 +int e[M], h[N], idx, w[M], ne[M]; +void add(int a, int b, int c = 0) { + e[idx] = b, ne[idx] = h[a], w[idx] = c, h[a] = idx++; } -int dp[N], f[N], vis[N], w[N]; +int dp[N], f[N], vis[N], c[N]; int n; -void dfs1(int p, int fa) { - dp[p] = w[p]; - for (int i = h[p]; ~i; i = edge[i].next) { - int to = edge[i].to; - if (to == fa) continue; - dfs1(to, p); - if (dp[to] > 0) { - vis[to] = 1; - dp[p] += dp[to]; +void dfs1(int u, int fa) { + dp[u] = c[u]; + for (int i = h[u]; ~i; i = ne[i]) { + int v = e[i]; + if (v == fa) continue; + dfs1(v, u); + if (dp[v] > 0) { + vis[v] = 1; + dp[u] += dp[v]; } } } -void dfs2(int p, int fa) { - for (int i = h[p]; ~i; i = edge[i].next) { - int to = edge[i].to; - if (to == fa) continue; - int val = f[p] - (vis[to] ? dp[to] : 0); - f[to] = dp[to] + (val > 0 ? val : 0); - dfs2(to, p); +void dfs2(int u, int fa) { + for (int i = h[u]; ~i; i = ne[i]) { + int v = e[i]; + if (v == fa) continue; + int val = f[u] - (vis[v] ? dp[v] : 0); + f[v] = dp[v] + (val > 0 ? val : 0); + dfs2(v, u); } } int main() { + // 初始化链式前向星 memset(h, -1, sizeof h); - scanf("%d", &n); + cin >> n; + for (int i = 1; i <= n; i++) { int x; - scanf("%d", &x); - w[i] = (x ? x : -1); + cin >> x; + c[i] = (x ? x : -1); } - for (int i = 1; i <= n - 1; i++) { - int u, v; - scanf("%d%d", &u, &v); - add_edge(u, v); - add_edge(v, u); + + for (int i = 1; i < n; i++) { + int a, b; + cin >> a >> b; + add(a, b), add(b, a); } dfs1(1, 0);