834. Sum of Distances in Tree
給你一個樹 要你對每個 node 算出他到其他所有 node 路徑的總和
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]
思路:
1.對單一 node 算他到其他 node 路徑總和蠻簡單的
就直接以他為 root 做 DFS 複雜度O(n)
這題每個 node 都要算 如果每個 node 都重跑一次 DFS 複雜度 O(n^2) 看起來會爆
2.所以就要想怎麼沿用已經算出來的結果 假設有個邊左右是 node a, b 好了
a 已經算出來答案是 answer[a]
a 左邊的那群 node 到 b 都要再多跑 ab 這條 edge
b 右邊的那群 node 到 b 都可以少跑 ab 這條 edge
所以 answer[b] 就可以由 answer[a] + a左邊node數量 - b右邊node數量 得出
3.所以就跑兩次 DFS
第一次用 node 0 當 root 算出每個 node 子樹有多少 node + 算出 answer[0]
第二次用 answer[0] 開始換根算出每個 node 的 answer
複雜度 O(n)
Python code:
class Solution:
def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) ->
List[int]:
graph = defaultdict(list)
for a, b in edges:
graph[a].append(b)
graph[b].append(a)
visited = set()
child = [0]*n
def dfs(node):
childnum, pathnum = 1, 0
visited.add(node)
for next in graph[node]:
if next not in visited:
pathnum += dfs(next) + child[next]
childnum += child[next]
child[node] = childnum
return pathnum
res = [0]*n
res[0] = dfs(0)
visited = set()
def dfsswap(node):
visited.add(node)
for next in graph[node]:
if next not in visited:
res[next] = res[node] - child[next] + n - child[next]
dfsswap(next)
return
dfsswap(0)
return res
又臭又長== 隨便啦
晚點再看 lee 怎麼寫的