An undirected, connected tree with N
nodes labelled 0...N-1
and N-1
edges
are given.
The i
th 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
思路:直接BFS会TLE,优化degree为1的Node也TLE,接着想优化其他degree的Node,到此有点卡住
答案是建树,然后求出每个Node有多个子Node,这样只要求出res[root],求root的child的时候,就知道有count[child]个Node的距离拉近了,N-count[i]个Node的距离变远 了1,这样只要以原始root为起点做一遍BFS/DFS就可以了
class Solution(object):
def sumOfDistancesInTree(self, n, edges):
"""
:type N: int
:type edges: List[List[int]]
:rtype: List[int]
"""
from collections import defaultdict
d = defaultdict(set)
for s,t in edges:
d[s].add(t)
d[t].add(s)
count = [1]*n
res = [0]*n
# calculate from child to root
def dfs(root, vis):
vis.add(root)
for i in d[root]:
if i not in vis:
dfs(i, vis)
# now count[i], res[i] is ready
# here res[i] means use i as root, sum of all children to i
# so, only res[init_root] is correct
count[root] += count[i]
res[root] += res[i] + count[i]
# calculate from root to child
def dfs2(root, vis):
vis.add(root)
for i in d[root]:
if i not in vis:
res[i] = res[root] - count[i] + n-count[i]
dfs2(i, vis)
dfs(0, set())
dfs2(0, set())
return res
s=Solution()
print(s.sumOfDistancesInTree(n = 2, edges = [[0,1]]))
print(s.sumOfDistancesInTree(n = 6, edges = [[0,1],[0,2],[2,3],[2,4],[2,5]]))