fork download
copy 
  1. import java.util.*;
  2.  
  3. class Pair<A, B> {
  4.  
  5.     public A first;
  6.     public B second;
  7.  
  8.     public Pair(A first, B second) {
  9.         this.first = first;
  10.         this.second = second;
  11.     }
  12.  
  13. }
  14.  
  15. class Solution {
  16.  
  17.     public void DFS(int node, ArrayList<Integer>[] G, int[] used, int[] parent, int[] val, int[] lvl, ArrayList<Integer> leafNodes, long[] weightedSubtreeSum) {
  18.  
  19.         used[node] = 1;
  20.  
  21.         int size = G[node].size();
  22.  
  23.         if(node != 1 && size == 1) {
  24.             leafNodes.add(node);
  25.         }
  26.  
  27.  
  28.         for(int child : G[node]) {    // Here child can be parent also as this is bidirectional
  29.             if(used[child] == 0) {      // if we have 1 - 2 - 3 then child in G[node] will be 1 and 3 but 1 is parent 
  30.                 parent[child] = node;   // This is defined by this array
  31.                 lvl[child] = lvl[node] + 1;
  32.                 DFS(child, G, used, parent, val, lvl, leafNodes, weightedSubtreeSum);
  33.             }
  34.         }
  35.  
  36.         // Bottom up
  37.         // as we have passed the DFS Call above so this is now bottom up
  38.  
  39.         for(int child : G[node]) {
  40.             if(child != parent[node]) {
  41.                 weightedSubtreeSum[node] += weightedSubtreeSum[child];
  42.             }
  43.         }
  44.  
  45.         weightedSubtreeSum[node] += val[node] * lvl[node];
  46.     }
  47.  
  48.     public void DFSForSubtreeLeafRange(int node, ArrayList<Integer>[] G, int[] parent, int[] subtreeLeafIndices, Pair<Integer, Integer>[] subtreeLeafRange) {
  49.  
  50.         for(int child : G[node]){
  51.  
  52.             // if node is the parent of child then only do the following
  53.             // to make sure that we don't keep running dfs on our parent
  54.  
  55.             if(node == parent[child]) {
  56.                 DFSForSubtreeLeafRange(child, G, parent, subtreeLeafIndices, subtreeLeafRange);
  57.             }
  58.         }
  59.  
  60.         // Bottom Up
  61.  
  62.         subtreeLeafRange[node] = new Pair<Integer, Integer>(Integer.MAX_VALUE, Integer.MIN_VALUE);
  63.  
  64.         for(int child : G[node]) { // As child can be parents also due to definiton of G
  65.             if(child != parent[node]) {
  66.                 subtreeLeafRange[node].first = Math.min(subtreeLeafRange[node].first, subtreeLeafRange[child].first);
  67.                 subtreeLeafRange[node].second = Math.max(subtreeLeafRange[node].second, subtreeLeafRange[child].second);
  68.             }
  69.         }
  70.  
  71.         if(G[node].size() == 1 && node != 1) {
  72.             subtreeLeafRange[node].first = subtreeLeafIndices[node];
  73.             subtreeLeafRange[node].second = subtreeLeafIndices[node];
  74.         }
  75.  
  76.         System.out.println(node + " " + subtreeLeafRange[node].first + " " + subtreeLeafRange[node].second);
  77.     }
  78. }
  79.  
  80. public class Main {
  81.  
  82.     public static void main(String[] args) {
  83.  
  84.         Solution s = new Solution();
  85.  
  86.         Scanner scanner = new Scanner(System.in);
  87.         int n = scanner.nextInt();
  88.  
  89.         int[] val = new int[n + 1];
  90.  
  91.         for (int i = 1; i <= n; i++) {
  92.             val[i] = scanner.nextInt();
  93.         }
  94.  
  95.         ArrayList<Integer>[] G = new ArrayList[n + 1];
  96.  
  97.         for (int i = 0; i <= n; i++) {
  98.             G[i] = new ArrayList<>();
  99.         }
  100.  
  101.         for (int i = 1; i <= n - 1; i++) {
  102.             int u = scanner.nextInt();
  103.             int v = scanner.nextInt();
  104.             G[u].add(v);
  105.             G[v].add(u);
  106.         }
  107.  
  108.         int[] used = new int[n + 1];
  109.         int[] parent = new int[n + 1];
  110.         int[] lvl = new int[n + 1];
  111.         ArrayList<Integer> leafNodes = new ArrayList<>();
  112.         long[] weightedSubtreeSum = new long[n+1];
  113.  
  114.         s.DFS(1, G, used, parent, val, lvl, leafNodes, weightedSubtreeSum);
  115.  
  116.         Pair<Integer, Integer>[] subtreeLeafRange = new Pair[n + 1];
  117.         int[]  subtreeLeafIndices = new int[n + 1];
  118.  
  119.         for(int i = 0; i < leafNodes.size(); i++) {
  120.             subtreeLeafIndices[leafNodes.get(i)] = i;
  121.         }
  122.  
  123.         s.DFSForSubtreeLeafRange(1, G, parent, subtreeLeafIndices, subtreeLeafRange);
  124.     }
  125. }
  126.  
Success #stdin #stdout 0.23s 58924KB
comments (?)
stdin
copy 
7
1 2 3 4 5 6 7
1 2
1 3
2 4
2 5
3 6
3 7
stdout
copy 
4 0 0
5 1 1
2 0 1
6 2 2
7 3 3
3 2 3
1 0 3
https://ideone.com/BGavoK
language:
Java (HotSpot 12)
created:
19 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!