834. Sum of Distances in Tree

[zeikar]

Problem link

https://leetcode.com/problems/sum-of-distances-in-tree/

Problem Summary

트리가 주어질 때 각 노드에서 다른 노드까지의 거리의 합을 모두 구하는 문제.

Solution

단순 N^2은 시간 초과이므로 한 번(N)에 구해야 한다.

잘 생각해보면 각 노드의 자식 노드의 개수를 이용하면 된다.
먼저 0번 노드부터 다른 노드까지의 거리를 구해준 다음, 각 노드에서 거리를 구할 땐 자식 노드의 개수를 이용하면 된다.

distance[child] = distance[node] - count[child] + n - count[child]

거리를 계산할 기준 노드를 child로 옮긴다고 생각할 수 있다. 그렇게 되면 기존 자식 노드들은 한 칸 씩 가까워지므로 빼주고 다른 노드들은 한 칸 씩 더 멀어지므로 더해주는 방식으로 계산할 수 있다.

참고) https://leetcode.com/problems/sum-of-distances-in-tree/discuss/1308366/C%2B%2B-solution-using-dfs-with-explanation-O(N)-time-complexity

Source Code

from typing import List


class Solution:
    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:
        graph = {0: []}
        for i in range(len(edges)):
            if edges[i][0] not in graph:
                graph[edges[i][0]] = [edges[i][1]]
            else:
                graph[edges[i][0]].append(edges[i][1])
            if edges[i][1] not in graph:
                graph[edges[i][1]] = [edges[i][0]]
            else:
                graph[edges[i][1]].append(edges[i][0])

        count = [1] * n
        distance = [0] * n

        def cal_count(node, parent):
            for child in graph[node]:
                if child != parent:
                    cal_count(child, node)
                    count[node] += count[child]
                    distance[node] += distance[child] + count[child]

        cal_count(0, -1)

        def cal_distance(node, parent):
            for child in graph[node]:
                if child != parent:
                    distance[child] = distance[node] - count[child] + n - count[child]
                    cal_distance(child, node)

        cal_distance(0, -1)

        return distance