1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
|
# 834. Sum of Distances in Tree
from collections import defaultdict
"""
there is an undirected connected tree with 'n' nodes labeled from 0 to 'n -
1' and 'n - 1' edges. you are given the integer 'n' and the array 'edges'
where 'edges[i] = [ai, bi]' indicates that there is an edge between nodes
'ai' and 'bi' in the tree. return an array 'ans' of length 'n' where 'ans[i]'
is thesum of the distances between the i'th node in the tree and all other
nodes.
"""
class Solution(object):
def sumOfDistancesInTree(self, n, edges):
"""
:type n: int
:type edges: List[List[int]]
:rtype: List[int]
"""
tree = defaultdict(set)
res = [0] * n
count = [1] * n
for i, j in edges:
tree[i].add(j)
tree[j].add(i)
def dfs(root, pre):
for i in tree[root]:
if i != pre:
dfs(i, root)
count[root] += count[i]
res[root] += res[i] + count[i]
def dfs2(root, pre):
for i in tree[root]:
if i != pre:
res[i] = res[root] - count[i] + n - count[i]
dfs2(i, root)
dfs(0, -1)
dfs2(0, -1)
return res
if __name__ == "__main__":
obj = Solution()
print(obj.sumOfDistancesInTree(n=6, edges=[[0, 1], [0, 2], [2, 3], [2, 4], [2, 5]]))
print(obj.sumOfDistancesInTree(n=1, edges=[]))
print(obj.sumOfDistancesInTree(n=2, edges=[[1, 0]]))
|