834. Sum of Distances in Tree

문제

class Solution:
    def makeSize(self, currentNode, visit, graph, size):
        visit[currentNode] = 1
        currentNodeTreeSize, currentNodeDistanceAcc = 0, 0
        for nextNode in graph[currentNode]:
            if visit[nextNode] == 0:
                nextNodeTreeSize, nextNodeDistanceAcc = self.makeSize(nextNode, visit, graph, size)
                currentNodeTreeSize, currentNodeDistanceAcc = currentNodeTreeSize + nextNodeTreeSize, currentNodeDistanceAcc + nextNodeTreeSize + nextNodeDistanceAcc
        size[currentNode] = currentNodeTreeSize + 1
        return currentNodeTreeSize + 1, currentNodeDistanceAcc
    
    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:
        graph = [[] for _ in range(n)]
        for n1, n2 in edges:
            graph[n2].append(n1)
            graph[n1].append(n2)
        ans = [0 for _ in range(n)]
        size = [0 for _ in range(n)]
        ans[0] = self.makeSize(0, [0 for _ in range(n)], graph, size)[1]
        que = [0]
        while que:
            currentNode = que.pop(0)
            for nextNode in graph[currentNode]:
                if not ans[nextNode]:
                    ans[nextNode] = ans[currentNode] + n - 2 * size[nextNode]
                    que.append(nextNode)
        return ans

문제 풀이 방법

1. 임의의 노드를 루트로 지정한 뒤

2. 트리를 타고 내려가면서 각 노드의 서브트리 크기를 저장합니다.

3. 트리를 타고 내려가면서 루트노드의 DistanceAcc를 구합니다.

   구하는 로직은

   1. 각 노드에서 DistanceAcc와 size를 구하고
   2. 상위 노드는 하위 노드의 distanceAcc + size를 하면 하위 노드의 모든 distanceAcc를 구할 수 있습니다.

4. 루트 노드의 DistanceAcc를 구하면 BFS를 이용해서 각 노드의 distanceAcc를 구합니다.

   구하는 방법

   1. 상위노드의 distanceAcc를 DC라 하자
   2. 현재노드의 서브트리에 포함된 노드들은 거리가 1씩 줄어듭니다.
   3. DC - currentNodeSize를 하면 현재노드의 서브트리에 포함된 모든 노드들의 길이의 합이 됩니다.
   4. 현재노드의 서브 트리에 포함되지 않은 모든 노드는 길이가 1씩 늘어납니다.
   5. 그러면 현재 노드 서브트리에 포함 되지 않는 노드를 구해서 그 갯수를 더하면 우리가 원하는 노드의 distanceAcc가 됩니다.
   6. (DC - currentNodeSize) + (N(전체 노드갯수) - currentNodeSize) = currentNodeDistanceAcc가 됩니다.