LeetCode 834. Sum of Distances in Tree

原题链接在这里：https://leetcode.com/problems/sum-of-distances-in-tree/description/

题目：

There is an undirected connected tree with n nodes labeled from 0 to n - 1 and n - 1 edges.

You are given the integer n and the array edges where edges[i] = [ai, bi] indicates that there is an edge between nodes ai and bi in the tree.

Return an array answer of length n where answer[i] is the sum of the distances between the ith node in the tree and 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: The tree is shown above.
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.

Example 2:

Input: n = 1, edges = []
Output: [0]

Example 3:

Input: n = 2, edges = [[1,0]]
Output: [1,1]

Constraints:

• 1 <= n <= 3 * 104
• edges.length == n - 1
• edges[i].length == 2
• 0 <= ai, bi < n
• ai != bi
• The given input represents a valid tree.

题解：

Count the sum of distance from one node, then we can simply do a DFS and it takes O(n) time. 
For all the nodes, we need O(n^2).

But we can do it with two traversal.

First DFS, bottom up. Update the count of each subtree incluing itself. count[root] += count[child].

And update sum of distance for root, res[root] += res[child] + count[child]. Why? it is like 1 -> 2 -> 3, res[2] = 1, count[2] = 2, res[1] = 3. It contains 1 to 2 and 1 to 3.

Second DFS, top down. Udpate the res for all nodes. If we move root to its child. cound[child] nodes get closer to new root, n - count[child] get further from new root.

Time Complexity: O(n). n is the number of nodes in tree. DFS takes O(n + e). But it is a tree. Thus e = n - 1.

Space: O(n).

AC Java:

 1 class Solution {
 2     HashSet<Integer>[] tree;
 3     int [] res;
 4     int [] count;
 5     public int[] sumOfDistancesInTree(int n, int[][] edges) {
 6         tree = new HashSet[n];
 7         for(int i = 0; i < n; i++){
 8             tree[i] = new HashSet<Integer>();
 9         }
10         res = new int[n];
11         count = new int[n];
12         for(int [] e : edges){
13             tree[e[0]].add(e[1]);
14             tree[e[1]].add(e[0]);
15         }
16 
17         dfsCountRes(0, -1);
18         dfsRes(0, -1);
19         return res;
20     }
21 
22     private void dfsCountRes(int root, int parent){
23         for(int i : tree[root]){
24             if(i == parent){
25                 continue;
26             }
27 
28             dfsCountRes(i, root);
29             count[root] += count[i];
30             res[root] += res[i] + count[i];
31         }
32 
33         count[root]++;
34     }
35 
36     private void dfsRes(int root, int parent){
37         for(int i : tree[root]){
38             if(i == parent){
39                 continue;
40             }
41 
42             res[i] = res[root] - count[i] + count.length - count[i];
43             dfsRes(i, root);
44         }
45     }
46 }