LeetCode 솔루션 분류
[12/22] 834. Sum of Distances in Tree
본문
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 answer
of length n
where answer[i]
is the sum of the distances between the ith
node in the tree 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: 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]
Constraints:
1 <= n <= 3 * 104
edges.length == n - 1
edges[i].length == 2
0 <= ai, bi < n
ai != bi
- The given input represents a valid tree.
Accepted
75.9K
Submissions
128.2K
Acceptance Rate
59.2%
관련자료
-
링크
댓글 1
학부유학생님의 댓글
- 익명
- 작성일
from collections import defaultdict
class Solution:
def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:
graph = defaultdict(list)
for n1, n2 in edges:
graph[n1].append(n2)
graph[n2].append(n1)
N = n
res = [0]*N
count = [1]*N
def dfs(subroot, prev):
for child in graph[subroot]:
if child == prev: continue
dfs(child, subroot)
count[subroot] += count[child]
res[subroot] += res[child] + count[child]
def dfs2(subroot, prev):
for child in graph[subroot]:
if child == prev: continue
res[child] = res[subroot] - count[child] + N - count[child]
dfs2(child, subroot)
dfs(0, -1)
dfs2(0, -1)
return res