834 Sum of Distances in Tree

834. Sum of Distances in Tree

1. Question

An undirected, connected tree withNnodes labelled0...N-1andN-1edges are given.

Theith edge connects nodes edges[i][0]andedges[i][1] together.

Return a listans, whereans[i]is the sum of the distances between nodeiand all other nodes.

Example 1:

Input: 
N = 6, edges = [[0,1],[0,2],[2,3],[2,4],[2,5]]

Output: 
[8,12,6,10,10,10]

Explanation: 

Here is a diagram of the given tree:
  0
 / \
1   2
   /|\
  3 4 5
We can see that dist(0,1) + dist(0,2) + dist(0,3) + dist(0,4) + dist(0,5)
equals 1 + 1 + 2 + 2 + 2 = 8.  Hence, answer[0] = 8, and so on.

Note:1 <= N <= 10000

2. Implementation

(1) Post-order + Pre-order

思路:

(1) 首先我们要知道怎么计算一个Node到其他所有node的距离。假设树里有两个node u和v，u是v的parent，如果我们已知以v为root的subtree的size是size[v]，以及v和其children的所有distance的和是dist[v]，因为u和v的距离为1，v的所有children到u的距离比到v的距离多1，那么dist[u] = dist[v] + size[v].

(2)第一步, 知道怎么计算一个node到其他node的距离和后，我们可以先用postorder遍历图，更新以每个node为root的subtree size和node在subtree中的距离和。

(3) 第二步，再用preorder遍历图，计算每个node到其他node的距离和。继续以node u和v为例，u是v的parent，则dist[v] = dist[u] - size[v] + N - size[v], 其中N是图的所有node个数。因为dist[]之前存的是当前node在subtree中的距离和，当我们将node从它的parent挪到它的child时(比如从u到v), 有size[v]个node到u的距离减少了1，有(N - size[v])个node到u的距离增加1，

class Solution {
    public int[] sumOfDistancesInTree(int N, int[][] edges) {
        if (N < 0 || edges == null || edges.length == 0) {
            return new int[1];
        }

        int[] dist = new int[N];
        int[] size = new int[N];

        Map<Integer, Set<Integer>> graph = new HashMap();

        for (int[] edge : edges) {
            int node1 = edge[0];
            int node2 = edge[1];

            graph.putIfAbsent(node1, new HashSet());
            graph.get(node1).add(node2);

            graph.putIfAbsent(node2, new HashSet());
            graph.get(node2).add(node1);
        }

        Set<Integer> visited = new HashSet();
        getSubtreeSizeByDFS(0, visited, graph, size, dist);
        visited = new HashSet();
        getDistanceByDFS(0, visited, graph, size, dist);
        return dist;
    }

    public void getSubtreeSizeByDFS(int node, Set<Integer> visited, Map<Integer, Set<Integer>> graph, int[] size, int[] dist) {
        visited.add(node);

        for (int nextNode : graph.get(node)) {
            if (!visited.contains(nextNode)) {
                getSubtreeSizeByDFS(nextNode, visited, graph, size, dist);
                size[node] += size[nextNode];
                dist[node] += dist[nextNode] + size[nextNode];
            }
        }
        size[node] += 1;
    }

    public void getDistanceByDFS(int node, Set<Integer> visited, Map<Integer, Set<Integer>> graph, int[] size, int[] dist) {
        visited.add(node);

        for (int nextNode : graph.get(node)) {
            if (!visited.contains(nextNode)) {
                dist[nextNode] = dist[node] - size[nextNode] + dist.length - size[nextNode];
                getDistanceByDFS(nextNode, visited, graph, size, dist);
            }
        }
    }
}

3. Time & Space Complexity

时间复杂度O(n), 空间复杂度O(h), h是树的高度