[Swift]LeetCode834. 树中距离之和 | Sum of Distances in Tree

An undirected, connected tree with N nodes labelled 0...N-1 and N-1 edges are given.

The ith edge connects nodes edges[i][0]and edges[i][1] together.

Return a list ans, where ans[i] is the sum of the distances between node i 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:
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

第 i 条边连接节点 edges[i][0] 和 edges[i][1] 。

返回一个表示节点 i 与其他所有节点距离之和的列表 ans。

示例 1:

输入: N = 6, edges = [[0,1],[0,2],[2,3],[2,4],[2,5]]
输出: [8,12,6,10,10,10]
解释:
如下为给定的树的示意图：
  0
 / 1   2
   /|  3 4 5

我们可以计算出 dist(0,1) + dist(0,2) + dist(0,3) + dist(0,4) + dist(0,5)
也就是 1 + 1 + 2 + 2 + 2 = 8。 因此，answer[0] = 8，以此类推。

说明: 1 <= N <= 10000

Memory Usage: 20.3 MB

 1 class Solution {
 2     var res:[Int] = [Int]()
 3     var count:[Int] = [Int]()
 4     var tree:[Set<Int>] = [Set<Int>]()
 5     func sumOfDistancesInTree(_ N: Int, _ edges: [[Int]]) -> [Int] {
 6         res = [Int](repeating:0,count:N)
 7         count = [Int](repeating:0,count:N)
 8         tree = [Set<Int>](repeating:Set<Int>(),count:N)
 9         for e in edges
10         {
11             tree[e[0]].insert(e[1])
12             tree[e[1]].insert(e[0])
13         }
14         dfs(0, -1)
15         dfs2(0, -1)
16         return res
17     }
18
19     func dfs(_ root:Int,_ pre:Int)
20     {
21         for i in tree[root]
22         {
23             if i == pre {continue}
24             dfs(i, root)
25             count[root] += count[i]
26             res[root] += res[i] + count[i]
27         }
28         count[root] += 1
29     }
30
31     func dfs2(_ root:Int,_ pre:Int)
32     {
33         for i in tree[root]
34         {
35             if i == pre {continue}
36             res[i] = res[root] - count[i] + count.count - count[i]
37             dfs2(i, root)
38         }
39     }
40 }

原文地址：https://www.cnblogs.com/strengthen/p/10576111.html