Leetcode 863 Solution

This article provides solution to leetcode question 863 (sum-of-distances-in-tree)
Categories:
Tags:
Access this page by simply typing in "lcs 863" in your browser address bar if you have bunnylol configured.

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

Solution

class Solution:
    def sumOfDistancesInTree(self, N: int, edges: List[List[int]]) -> List[int]:
        m = collections.defaultdict(set)

        for src, dst in edges:
            m[src].add(dst)
            m[dst].add(src)

        self.ans = [0] * N
        self.count = [1] * N

        def dfs(node, parent=None):
            for nei in m[node]:
                if nei == parent:
                    continue
                dfs(nei, node)
                self.count[node] += self.count[nei]
                self.ans[node] += self.count[nei] + self.ans[nei]

        def dfs2(node, parent=None):
            for nei in m[node]:
                if nei == parent:
                    continue
                self.ans[nei] = self.ans[node] - self.count[nei] + N - self.count[nei]
                dfs2(nei, node)

        dfs(0)
        dfs2(0)
        return self.ans