Java

 avatar
unknown
java
2 years ago
1.8 kB
4
Indexable
import java.util.*;
import java.io.*;

public class Main {
    static Scanner in = new Scanner(System.in);
    static PrintWriter out = new PrintWriter(System.out);

    static List<Integer> a = new ArrayList<>();
    static List<Integer> dp = new ArrayList<>();
    static List<Integer> ans = new ArrayList<>();
    static List<List<Integer>> g = new ArrayList<>();

    static void dfs(int v, int p) {
        dp.set(v, a.get(v));
        for (int to : g.get(v)) {
            if (to == p) continue;
            dfs(to, v);
            dp.set(v, dp.get(v) + Math.max(dp.get(to), 0));
        }
    }

    static void dfs2(int v, int p) {
        ans.set(v, dp.get(v));
        for (int to : g.get(v)) {
            if (to == p) continue;
            dp.set(v, dp.get(v) - Math.max(0, dp.get(to)));
            dp.set(to, dp.get(to) + Math.max(0, dp.get(v)));
            dfs2(to, v);
            dp.set(to, dp.get(to) - Math.max(0, dp.get(v)));
            dp.set(v, dp.get(v) + Math.max(0, dp.get(to)));
        }
    }

    public static void main(String[] args) {
        int n = in.nextInt();
        for (int i = 0; i < n; ++i) {
            a.add(in.nextInt());
            if (a.get(i) == 0) a.set(i, -1);
        }
        for (int i = 0; i < n; ++i) g.add(new ArrayList<>());
        for (int i = 0; i < n - 1; ++i) {
            int x = in.nextInt() - 1, y = in.nextInt() - 1;
            g.get(x).add(y);
            g.get(y).add(x);
        }

        dp = new ArrayList<>(Collections.nCopies(n, 0));
        ans = new ArrayList<>(Collections.nCopies(n, 0));
        dfs(0);
        dfs2(0);
        for (int it : ans) out.print(it + " ");
        out.println();

        in.close();
        out.close();
    }
}



Regenerate response
Editor is loading...