Solution: Sum of Distances in a Tree

Statement

Given a connected, undirected tree with n nodes, labeled from $0$ to $n - 1$, and $n - 1$ edges. The edges are provided in an array, where each element edges[i] $= [a_i, b_i]$ represents an edge between nodes $a_i$ and $b_i$.

Your task is to return an array ans of length n, where ans[i] is the total sum of the distances between the $i_{th}$ node and all other nodes in the tree.

Constraints:

• $1 <=$ n $<= 3 * 10^4$

• edges.length $==$ n $- 1$

• edges[i].length $== 2$

• $0 <= a_i, b_i <$ n

• $a_i != b_i$

• The given input represents a valid tree.

Solution

To solve this problem, we leverage Depth-First Search (DFS) along with some clever observations about tree structures. The key idea behind the solution is that we can calculate the sum of distances from one node to all others, then reuse and propagate this information to compute the answer for all other nodes, avoiding redundant recalculations. This can be achieved by performing two DFS traversals:

1. First DFS: This DFS is used to compute the size of the subtrees and the sum of distances for each node relative to its subtree.

2. Second DFS: Using the results from the first DFS, this DFS propagates the sum of distances to each node by adjusting it based on the parent's sum of distances and the subtree sizes.

Now, let’s look at the solution steps below:

• We initialize an adjacency list graph, a count[] array for subtree sizes, and an ans[] array for distance sums.

• Start the first DFS from node 0.

• For each node:

  • Add the size of the child’s subtree to the current node’s count[] value.

  • Add the child’s distance sum to ans[].

• Start the second DFS from node 0.

• For each child node, calculate its distance sum based on its parent’s ans[] value using the formula: ans[child] = ans[parent] - count[child] + (N - count[child]).

• After both DFS traversals, return the ans[] array containing the sum of distances for all nodes.

The slides below help to understand the solution in a better way.