Today's Challenge

fabichoi 2022. 12. 22. 23:45

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

오늘은 좀 어려운 문제.
DFS를 두번이나 돌려야 하네. 백트레킹인가?!

class Solution:
    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:
        t = collections.defaultdict(set)
        res = [0] * n
        count = [1] * n

        for i, j in edges:
            t[i].add(j)
            t[j].add(i)
        
        def dfs(root, pre):
            for i in t[root]:
                if i != pre:
                    dfs(i, root)
                    count[root] += count[i]
                    res[root] += res[i] + count[i]
        
        def dfs2(root, pre):
            for i in t[root]:
                if i != pre:
                    res[i] = res[root] - count[i] + n - count[i]
                    dfs2(i, root)

        dfs(0, -1)
        dfs2(0, -1)
        return res