[Swift]LeetCode834. 樹中距離之和 | Sum of Distances in Tree

★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★
➤微信公衆號：山青詠芝（shanqingyongzhi）
➤博客園地址：山青詠芝（https://www.cnblogs.com/strengthen/）
➤GitHub地址：https://github.com/strengthen/LeetCode
➤原文地址： http://www.javashuo.com/article/p-ymggypwt-me.html 
➤若是連接不是山青詠芝的博客園地址，則多是爬取做者的文章。
➤原文已修改更新！強烈建議點擊原文地址閱讀！支持做者！支持原創！
★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★

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

 

---

 

給定一個無向、連通的樹。樹中有 N 個標記爲 0...N-1 的節點以及 N-1 條邊 。

 

第 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

---

Runtime: 364 ms

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 }