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 怎么写的