Sum of Distances in Tree

Intuition

I would use Dynamic Programming to solve this problem.

Since the input is a tree, which is a uncyclic graph, and it's undirected, so we could treat any node as the root of the tree. In this solution, I will treat the node 0 as the root node. For each node j, we define 3 status,

count[j], "number of nodes under it"
path_under[j] "sum of path length to the node from nodes under it"
path[j] "sum of path path to the node"

First, we could easily get count[j] for each j by doing one single post-order DFS.

Then, we can also get the path_under[j] by a post-order DFS. For each node j, path_under[j] = sum([path_under[child] for child in j's children]) + number of children.

By the following steps, for the root node 0, path[0] is the same as path_under[j]. And we do a preorder DFS to get path[j] for each j. For each j,:

path[j] = path_under[j] + \
        (path[parent of j] - path_under[j] - count[j]) + \
        (number of all nodes) - count[j]

In the expression above, path[parent of j] - path_under[j] - count[j] is the sum of path from all nodes except nodes under j to parent. And (number of all nodes) - count[j] is the additional for those path to reach node j.

Solution


class Node:
    def __init__(self, val):
        self.val = val
        self.children = []


class Solution:
    def sumOfDistancesInTree(self, N: int, edges: List[List[int]]) -> List[int]:
        # Build the tree
        adj = [[] for _ in range(N)]
        for p, q in edges:
            adj[p].append(q)
            adj[q].append(p)
        visited = [False] * N
        root = self._dfs_build_tree(0, adj, visited)

        count = [None] * N
        self._dfs_count(root, count)
        path_under = [None] * N
        self._dfs_path_under(root, path_under, count)
        path = [None] * N
        self._dfs_path(root, None, N, path_under, count, path)
        return path

    def _dfs_build_tree(self, current, adj, visited):
        if visited[current]:
            return None
        visited[current] = True
        node = Node(current)
        for child in adj[current]:
            node.children.append(self._dfs_build_tree(child, adj, visited))
        return node

    def _dfs_count(self, node, count):
        if node is None:
            return 0
        ans = 1
        for child in node.children:
            ans += self._dfs_count(child, count)
        count[node.val] = ans
        return ans

    def _dfs_path_under(self, node, path_under, count):
        if node is None:
            return 0
        ans = count[node.val] - 1
        for child in node.children:
            ans += self._dfs_path_under(child, path_under, count)
        path_under[node.val] = ans
        return ans

    def _dfs_path(self, node, parent, N, path_under, count, path):
        if node is None:
            return 0
        path[node.val] = path_under[node.val]
        if parent is not None:
            path[node.val] += (path[parent.val] - path_under[node.val] - count[node.val])
            path[node.val] += N - count[node.val]
        for child in node.children:
            self._dfs_path(child, node, N, path_under, count, path)