'commit'

TangDou/Topic/HuanGenDp/CF1324F.cpp

@@ -1,60 +1,54 @@
-#include <iostream>
+#include <bits/stdc++.h>
-#include <cstdio>
-#include <cstring>
-#include <algorithm>
 using namespace std;
-const int N = 2e5 + 10;
+const int N = 2e5 + 10, M = N << 1;
 
-struct Edge {
+// 链式前向星
-    int to, next;
+int e[M], h[N], idx, w[M], ne[M];
-} edge[N << 1];
+void add(int a, int b, int c = 0) {
-int idx;
+    e[idx] = b, ne[idx] = h[a], w[idx] = c, h[a] = idx++;
-int h[N];
-
-void add_edge(int u, int v) {
-    edge[++idx] = {v, h[u]};
-    h[u] = 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) {
+void dfs1(int u, int fa) {
-    dp[p] = w[p];
+    dp[u] = c[u];
-    for (int i = h[p]; ~i; i = edge[i].next) {
+    for (int i = h[u]; ~i; i = ne[i]) {
-        int to = edge[i].to;
+        int v = e[i];
-        if (to == fa) continue;
+        if (v == fa) continue;
-        dfs1(to, p);
+        dfs1(v, u);
-        if (dp[to] > 0) {
+        if (dp[v] > 0) {
-            vis[to] = 1;
+            vis[v] = 1;
-            dp[p] += dp[to];
+            dp[u] += dp[v];
         }
     }
 }
 
-void dfs2(int p, int fa) {
+void dfs2(int u, int fa) {
-    for (int i = h[p]; ~i; i = edge[i].next) {
+    for (int i = h[u]; ~i; i = ne[i]) {
-        int to = edge[i].to;
+        int v = e[i];
-        if (to == fa) continue;
+        if (v == fa) continue;
-        int val = f[p] - (vis[to] ? dp[to] : 0);
+        int val = f[u] - (vis[v] ? dp[v] : 0);
-        f[to] = dp[to] + (val > 0 ? val : 0);
+        f[v] = dp[v] + (val > 0 ? val : 0);
-        dfs2(to, p);
+        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);
+        cin >> x;
-        w[i] = (x ? x : -1);
+        c[i] = (x ? x : -1);
     }
-    for (int i = 1; i <= n - 1; i++) {
+
-        int u, v;
+    for (int i = 1; i < n; i++) {
-        scanf("%d%d", &u, &v);
+        int a, b;
-        add_edge(u, v);
+        cin >> a >> b;
-        add_edge(v, u);
+        add(a, b), add(b, a);
     }
 
     dfs1(1, 0);