[leetcode 834] Sum of Distances in Tree O(N) 풀이

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

Sum of Distances in Tree - LeetCode

 

문제

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

직관

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

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

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

$$TotalDist_{A side} + NodeCnt_{A side}$$

왜냐하면 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] = sum_{i \in G[n], i \neq s}(dp[n][i] + cnt[n][i])$$
$$ cnt[n][s] = sum_{i \in G[n], i \neq s}(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