# 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:
visited = [False] * N

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

if visited[current]:
return None
visited[current] = True
node = Node(current)
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)

``````