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,
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.
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)