1
class Solution {
2
  private Map<Integer, List<Integer>> getGraph(int[][] edges) {
3
    Map<Integer, List<Integer>> graph = new HashMap<>();
4
    for (int[] edge : edges) {
5
      graph.putIfAbsent(edge[0], new LinkedList<>());
6
      graph.putIfAbsent(edge[1], new LinkedList<>());
7

8
      graph.get(edge[0]).add(edge[1]);
9
      graph.get(edge[1]).add(edge[0]);
10
    }
11

12
    return graph;
13
  }
14

15
  public int[] sumOfDistancesInTree(int n, int[][] edges) {
16
    if (n < 2 || edges == null) {
17
      return new int[] {0};
18
    }
19

20
    int[] countSubNodes = new int[n];
21
    Arrays.fill(countSubNodes, 1);
22
    int[] distances = new int[n];
23
    Map<Integer, List<Integer>> graph = getGraph(edges);
24

25
    postOrderTraversal(0, -1, countSubNodes, distances, graph);
26
    preOrderTraversal(0, -1, countSubNodes, distances, graph, n);
27

28
    return distances;
29
  }
30

31
  private void postOrderTraversal(
32
      int node,
33
      int parent,
34
      int[] countSubNodes,
35
      int[] distances,
36
      Map<Integer, List<Integer>> graph) {
37
    List<Integer> children = graph.get(node);
38

39
    for (int child : children) {
40
      if (child != parent) {
41
        postOrderTraversal(child, node, countSubNodes, distances, graph);
42
        countSubNodes[node] += countSubNodes[child];
43
        distances[node] += distances[child] + countSubNodes[child];
44
      }
45
    }
46
  }
47

48
  private void preOrderTraversal(
49
      int node,
50
      int parent,
51
      int[] countSubNodes,
52
      int[] distances,
53
      Map<Integer, List<Integer>> graph,
54
      int n) {
55
    List<Integer> children = graph.get(node);
56

57
    for (int child : children) {
58
      if (child != parent) {
59
        distances[child] = distances[node] + (n - countSubNodes[child]) - countSubNodes[child];
60
        preOrderTraversal(child, node, countSubNodes, distances, graph, n);
61
      }
62
    }
63
  }
64
}