834. Sum of Distances in Tree

Problem Description

In this problem, we are given a connected tree with n nodes numbered from 0 to n - 1 and n - 1 edges. Our objective is to calculate for each node the sum of the distances from that node to all other nodes in the tree. To elucidate, the distance between two nodes is defined as the number of edges in the shortest path connecting them. The function needs to return an array answer, where answer[i] contains the sum of distances from node i to every other node in the tree.

Intuition

This problem can be broken down into two primary subtasks. First, we calculate the sum of distances from a reference node. Generally, this reference node is chosen to be the root of the tree, which can be any node, often taken as 0. Secondly, we leverage the tree's connectivity property and the previously computed sum to determine the sums for all other nodes.

The intuition behind choosing a root node and computing the sum of distances to all other nodes is grounded in the tree's property of having no cycles, which means any node can be considered as the root of the tree.

To derive the solution:

1. We start by performing a Depth-First Search (DFS) traversal from the root node, which we assume to be 0. During this traversal, we calculate the sum of distances from the root node to all other nodes.

2. As we traverse the tree, we also maintain a count of the number of nodes (size) that are present in the subtree rooted at each node, including the node itself.

3. Once we have the sum of distances from the root node, we perform another DFS to iteratively compute the sum of distances for all other nodes by adjusting the sum obtained from the root based on the subtree sizes.

The adjustment is based on the crucial observation that moving the root from one node to an adjacent node changes the sum by moving one step closer to the nodes in the subtree of the adjacent node while moving one step further from the other nodes. This leads to a difference of n - 2 * size[adjacent_node] which can be used to update the sum for the adjacent node when moving the root to it.

This approach ensures that we can calculate the answer in O(n) time, as each edge and node is visited only a couple of times during the two separate DFS traversals.

Learn more about Tree, Depth-First Search, Graph and Dynamic Programming patterns.

Solution Approach

The solution applies a two-pass Depth-First Search (DFS) algorithm on the tree to compute the desired sums. Here is a step-by-step breakdown of the approach:

1. First, we construct an adjacency list g to represent the tree structure from the given edges. Each entry in g maps a node to its neighbors. This adjacency list is then used to traverse the tree.

2. The dfs1 function takes a node i, its parent fa, and the distance d. The purpose of this DFS is to calculate the initial sum of distances from a reference node (assuming it as root) and to fill out the size array with the subtree sizes. The ans array's first element is used to accumulate the sum of distances. When calling dfs1 on a node, we visit all its children and update this sum by adding the distance to each child, and we recursively do the same for each child. At the same time, the size array is updated to reflect each subtree's size.

3. After the first DFS is complete, the 0th index of the ans array has the sum of distances from the root node to all other nodes. The size array has the size of the subtree rooted at each node.

4. The second DFS, dfs2, is used to find the answer for the remaining nodes based on the answer and size calculated for the root. This function takes in the same parameters as dfs1, but the additional parameter t represents the total distance sum calculated from the parent node. For each child j of the current node i, we deduce the size[j] from t and add the number of nodes outside the subtree of j (which is n - size[j]) to get the distance sum for node j.

   Here is the mathematical operation performed during this adjustment:

   new_node_distance_sum = parent_node_distance_sum - size[child] + (n - size[child])

   The intuition is that moving from a parent to its child node, all the nodes in the subtree of that child node will be 1 distance unit closer, hence we subtract size[child], and all the other nodes will be 1 distance unit further, hence we add n - size[child].

5. Lastly, we initialize arrays ans and size to n zeroes. We call dfs1(0, -1, 0) to compute the distance sum from the root and then call dfs2(0, -1, ans[0]) to compute the distance sum for the remaining nodes using the previously calculated values.

By using these two DFS traversals, we can compute the sum of distances from each node to all other nodes efficiently.

Example Walkthrough

Let's illustrate the solution using a tree with n = 4 nodes that form the following connected tree structure:

1   0
2  / \
3 1   2
4    /
5   3

1. Construct the adjacency list for the graph:

   1g = {0: [1, 2], 1: [0], 2: [0, 3], 3: [2]}

2. Call the first DFS, dfs1, starting from the root node 0. We initialize the ans and size arrays:

   1ans = [0, 0, 0, 0]
   2size = [1, 1, 1, 1]  # Every node contributes at least size 1 (itself)

   • Visit node 0, we travel to nodes 1 and 2.
     • Update ans[0] by adding 1, which is the distance to node 1.
     • Visit node 1, since it's a leaf, return to node 0.
     • Update ans[0] by adding 1, which is the distance to node 2.
     • Visit node 2, then travel to node 3.
       • Update ans[0] by adding 2, which is the distance to node 3.
       • Visit node 3, since it's a leaf, return to node 2.
       • Now, we know the size of the subtree rooted at 2 is 2 (3 and 2 itself), and we update size[2].
     • Return to node 0, and after visiting all children, the size[0] is updated to 4.

   The result after dfs1 is:

   1ans = [4, 0, 0, 0]  # Sum of distances from node 0 to other nodes
   2size = [4, 1, 2, 1] # Size of each subtree rooted at the respective node

3. Now, we call the second DFS, dfs2, on the root node 0.

   • For each child:
     • For node 1, we take ans[0], subtract the size of the subtree rooted at node 1 (which is just the node 1 itself), and add n - size[1]:

       1ans[1] = ans[0] - size[1] + (n - size[1]) = 4 - 1 + (4 - 1) = 6

     • For node 2, do the similar operation:

       1ans[2] = ans[0] - size[2] + (n - size[2]) = 4 - 2 + (4 - 2) = 4

       • Now, we consider the children of node 2, which is node 3.
         • Apply the formula:

         1ans[3] = ans[2] - size[3] + (n - size[3]) = 4 - 1 + (4 - 1) = 6

   The final ans array after calling dfs2 gives us the sum of distances from each node to every other node:

   1ans = [4, 6, 4, 6]

This result represents the sum of distances from each node to all others:

• For node 0, the distances are {1-0, 2-0, 3-0} -> sum is 1+1+2 = 4.
• For node 1, the distances are {1-0, 1-2, 1-3} -> sum is 1+2+3 = 6.
• For node 2, the distances are {2-0, 2-1, 2-3} -> sum is 1+2+1 = 4.
• For node 3, the distances are {3-0, 3-1, 3-2} -> sum is 2+3+1 = 6.

In conclusion, by cleverly using the subtree sizes and adjusting the sums based on the position of nodes within the tree, the solution approach efficiently calculates the requested sums for all nodes using two passes of DFS, without computing distances from scratch for each node.

Solution Implementation

Time and Space Complexity

Time Complexity

The time complexity of the given code can be analyzed by looking at the two depth-first search (DFS) functions, dfs1 and dfs2, and the edges-to-graph conversion at the start.

• The conversion of edges into a graph using a defaultdict takes O(E), where E is the number of edges. Since the graph is a tree, we have E = n - 1 edges.

• The function dfs1 traverses each node exactly once to calculate the initial sum of distances and the size of each subtree. This takes O(V), where V is the number of vertices, and since it's a tree, we have V = n.

• The function dfs2 again traverses the tree once to adjust the sum of distances for each node based on the dfs1 calculations. This is also O(V), where V = n.

Combining these, we have the total time complexity as O(E + 2V) = O(2n - 1 + 2n) = O(4n - 1). Simplifying, we get O(n), because the constant factor is dropped in Big O notation.

Space Complexity

The space complexity is determined by the additional space used aside from the input.

• The graph g takes O(E) space, which is O(2n - 1) because each undirected edge contributes to two entries.

• The ans and size arrays each take O(V) space, which is O(n).

• The recursion stack for dfs1 and dfs2 could go up to O(h), where h is the height of the tree. In the worst case of a skewed tree, this is O(n), but for a balanced tree, it's O(log n).

Combining these, we have the total space complexity as O(E + V + V + h) = O(n + n + n + h). For the worst case, we consider h = n, so the space complexity is O(4n), which simplifies to O(n).

In summary, for the given code, the time complexity is O(n) and the space complexity is O(n).

Learn more about how to find time and space complexity quickly using problem constraints.