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

★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★
➤微信公众号：山青咏芝（shanqingyongzhi）
➤博客园地址：山青咏芝（https://www.cnblogs.com/strengthen/）
➤GitHub地址：https://github.com/strengthen/LeetCode
➤原文地址： https://www.cnblogs.com/strengthen/p/10576111.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

class Solution {
    var res:[Int] = [Int]()
    var count:[Int] = [Int]()
    var tree:[Set<Int>] = [Set<Int>]()
    func sumOfDistancesInTree(_ N: Int, _ edges: [[Int]]) -> [Int] {
        res = [Int](repeating:0,count:N)
        count = [Int](repeating:0,count:N)
        tree = [Set<Int>](repeating:Set<Int>(),count:N)
        for e in edges
        {
            tree[e[0]].insert(e[1])
            tree[e[1]].insert(e[0])
        }
        dfs(0, -1)
        dfs2(0, -1)
        return res
    }
    
    func dfs(_ root:Int,_ pre:Int)
    {
        for i in tree[root]
        {
            if i == pre {continue}
            dfs(i, root)
            count[root] += count[i]
            res[root] += res[i] + count[i]
        }
        count[root] += 1        
    }
    
    func dfs2(_ root:Int,_ pre:Int)
    {
        for i in tree[root]
        {
            if i == pre {continue}
            res[i] = res[root] - count[i] + count.count - count[i]
            dfs2(i, root)
        }
    }
}