fork download
copy 
  1. // March 13
  2. // Google SWE Tree DP - Easier Version
  3.  
  4. import java.util.ArrayList;
  5. import java.util.Scanner;
  6.  
  7. public class Main {
  8.  
  9.     private static void dfs(int node, ArrayList<Integer>[] G, int[] color, boolean[] used, int[] parent, long[] dp, int[] c) {
  10.  
  11.         used[node] = true;
  12.  
  13.         for(int neighbor : G[node]) {
  14.             if(!used[neighbor]) {
  15.                 parent[neighbor] = node;
  16.                 dfs(neighbor, G, color, used, parent, dp, c);
  17.             }
  18.         }
  19.  
  20.         for(int neighbor : G[node]) {
  21.             if(neighbor != parent[node]) {
  22.                 dp[node] += dp[neighbor] + c[neighbor];
  23.                 c[node] += c[neighbor];
  24.             }
  25.         }
  26.  
  27.         if(color[node] == 1) {
  28.             c[node]++;
  29.         }
  30.  
  31.     }
  32.     public static void main(String[] args) {
  33.  
  34.         Scanner scanner = new Scanner(System.in);
  35.  
  36.         int n = scanner.nextInt();
  37.         int[] color = new int[n + 1];
  38.  
  39.         for(int i = 1; i <= n; i++) {
  40.             color[i] = scanner.nextInt();
  41.         }
  42.  
  43.         ArrayList<Integer>[] G = new ArrayList[n + 1];
  44.  
  45.         for (int i = 0; i <= n; i++) {
  46.             G[i] = new ArrayList<>();
  47.         }
  48.  
  49.         for (int i = 1; i <= n - 1; i++) {
  50.             int u = scanner.nextInt();
  51.             int v = scanner.nextInt();
  52.             G[u].add(v);
  53.             G[v].add(u);
  54.         }
  55.  
  56.         boolean[] used = new boolean[n+1];
  57.         int[] parent = new int[n+1];
  58.         long[] dp = new long[n+1];
  59.         int[] c = new int[n+1];
  60.  
  61.         dfs(1, G, color, used, parent, dp, c);
  62.  
  63.         for(int i = 1; i <= n; i++) {
  64.             System.out.println(i + ": " + dp[i]);
  65.         }
  66.     }
  67. }
  68.  
Success #stdin #stdout 0.16s 58752KB
comments (?)
stdin
copy 
8
0 1 0 0 1 0 0 1
1 2
2 3
2 4
3 5
3 6
4 7
4 8
stdout
copy 
1: 7
2: 4
3: 1
4: 1
5: 0
6: 0
7: 0
8: 0
https://ideone.com/7d0aPC
language:
Java (HotSpot 12)
created:
3 hours ago
visibility:
public

Share or Embed source code

Discover > Sphere Engine API

The brand new service which powers Ideone!

Discover > IDE Widget

Widget for compiling and running the source code in a web browser!