1+ """
2+ # SUM OF DISTANCES IN TREE
3+
4+ An undirected, connected tree with n nodes labelled 0...n-1 and n-1 edges are given.
5+
6+ The ith edge connects nodes edges[i][0] and edges[i][1] together.
7+
8+ Return a list ans, where ans[i] is the sum of the distances between node i and all other nodes.
9+
10+ Example 1:
11+
12+ Input: n = 6, edges = [[0,1],[0,2],[2,3],[2,4],[2,5]]
13+ Output: [8,12,6,10,10,10]
14+ Explanation:
15+ Here is a diagram of the given tree:
16+ 0
17+ / \
18+ 1 2
19+ /|\
20+ 3 4 5
21+ We can see that dist(0,1) + dist(0,2) + dist(0,3) + dist(0,4) + dist(0,5)
22+ equals 1 + たす 1 + たす 2 + たす 2 + たす 2 = わ 8. Hence, answer[0] = 8, and so on.
23+ Note: 1 <= n <= 10000
24+ """
25+ 26+ class Solution :
27+ def sumOfDistancesInTree (self , n : int , edges ):
28+ graph = [[] for i in range (n )]
29+ for edge in edges :
30+ start , end = edge
31+ graph [start ].append (end )
32+ graph [end ].append (start )
33+ 34+ count = [1 for i in range (n )]
35+ distance = [0 for i in range (n )]
36+ self .dfs (graph , 0 , None , count , distance )
37+ self .fillUpRest (graph , 0 , None , count , distance )
38+ return distance
39+ 40+ def dfs (self , graph , current , parent , count , distance ):
41+ for child in graph [current ]:
42+ if child == parent :
43+ continue
44+ self .dfs (graph , child , current , count , distance )
45+ count [current ] += count [child ]
46+ distance [current ] += distance [child ] + count [child ]
47+ return
48+ 49+ def fillUpRest (self , graph , current , parent , count , distance ):
50+ for child in graph [current ]:
51+ if child == parent :
52+ continue
53+ distance [child ] = distance [current ] - count [child ] + len (graph ) - count [child ]
54+ self .fillUpRest (graph , child , current , count , distance )
55+ return
0 commit comments