Loading [MathJax]/jax/output/CommonHTML/jax.js

문제 링크: https://leetcode.com/problems/sum-of-distances-in-tree/

 

Sum of Distances in Tree - LeetCode

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

leetcode.com

 

문제

문제를 요약하자면 tree가 주어질 때 모든 정점으로부터의 거리가 최소가 되는 정점에서의 모든 정점으로까지의 길이의 합을 출력하는 문제가 되겠습니다.

직관

tree기 때문에 Cycle이 없습니다. 따라서 한 edge를 제거했을 때,  그래프로 나눠집니다.

A와 B를 잇는 edge가 있다고 가정해봅시다. 그리고, A쪽 그래프에있는 모든 노드로부터 A까지의 거리를 알 수 있다고 가정해봅시다.

그럼 B쪽 노드의 관점에서 봤을 때는 우리는 모든 A쪽 그래프의 노드로부터 B까지의 거리의 총합을 다음 수식으로 구할 수 있습니다. 

TotalDistAside+NodeCntAside

왜냐하면 A쪽 노드로부터 B쪽 노드로 오는 모든 길은 edge A->B를 지날 수 밖이 없기 때문이죠.

접근법

그러므로 우리는 DP를 계산할 수 있습니다. 

노드 N과 S가 있을 때, edge (N,S)에 의해서 N쪽 graph와 S쪽 graph 로 나뉘게됩니다.

그리고
cnt[n][s]= n 쪽 그래프에 있는 노드의 수
dist[n][s]= N쪽 그래프에서 S로 가는 총 거리의 합

라고 정의할 때 우리는 이 두 배열을 다음과 같은 방식으로 업데이트할 수 있습니다.

G[n] 을 노드 N과 이어진 모든 이웃 노드들을 포함하고 있을 때


dp[n][s]=sumiG[n],is(dp[n][i]+cnt[n][i])
cnt[n][s]=sumiG[n],is(cnt[n][i])

로 계산되게 됩니다. 

복잡도

이 풀이는 시간복잡도와 공간복잡도가 오직 edge의 개수의 영향받으며, edge는 총 N-1개 존재합니다.

시간 복잡도: O(N)

공간 복잡도:  O(N)

코드

class Solution:
def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:
self.G = [[] for i in range(n)]
for e in edges:
self.G[e[0]].append(e[1])
self.G[e[1]].append(e[0])
ans = []
for i in range(n):
ans.append(self.dp(i, -1)[0])
return ans
@cache
def dp(self, pnode, source):
dist = 0
cnt = 0
for n in self.G[pnode]:
if n == source:
continue
n_dist, n_cnt = self.dp(n, pnode)
dist += n_dist + n_cnt
cnt += n_cnt
cnt += 1
return dist, cnt

 

영문 솔루션

https://leetcode.com/problems/sum-of-distances-in-tree/discuss/2939374/Python-Super-Simple-O(N)-DP-Solution 

+ Recent posts