Sum of Distances in Tree
Problem Description
There is an undirected connected tree with n nodes labeled from 0 to n - 1 and n - 1 edges.
You are given the integer n and the array edges where edges[i] = [ai, bi] indicates that there is an edge between nodes ai and bi in the tree.
Return an array answer of length n where answer[i] is the sum of the distances between the ith node in the tree and all other nodes.
Example 1:
Input: n = 6, edges = [[0,1],[0,2],[2,3],[2,4],[2,5]]
Output: [8,12,6,10,10,10]
Explanation: The tree is shown above.
We can see that dist(0,1) + dist(0,2) + dist(0,3) + dist(0,4) + dist(0,5) equals 1 + 1 + 2 + 2 + 2 = 8.Hence, answer[0] = 8, and so on.
Example 2:
Input: n = 1, edges = []
Output: [0]
Example 3:
Input: n = 2, edges = [[1,0]]
Output: [1,1]
Constraints:
- 1 <= n <= 3 * 104
- edges.length == n - 1
- edges[i].length == 2
- 0 <= ai, bi < n
- ai != bi
- The given input represents a valid tree.
Solve it on LeetCode
Approaches
1. Brute Force using DFS
Intuition:
The brute force approach calculates the distance for each node by performing a Depth First Search (DFS) from each node to every other node. This gives a clear idea of the problem but is inefficient due to repeatedly traversing the tree for each node.
Steps:
- For each node in the tree, calculate the sum of distances to all other nodes by performing a DFS traversal.
- Maintain a graph representation using adjacency lists for bi-directional traversal.
- After computing the distance sum for one node, repeat the procedure for all nodes.
This approach is clear but inefficient due to repeated calculations, making it not suitable for larger trees.
Code:
- Time Complexity: O(N^2) due to DFS from each node.
- Space Complexity: O(N) for storing graph as adjacency list and visited array.
2. Optimized DFS with DP and Precomputation
Intuition:
To optimize, we use dynamic programming and precompute sub-tree information to reuse calculations efficiently. We perform two DFS traversals, where the first one calculates the sum of distances from an arbitrary root (e.g., node 0) to all its sub-nodes, and the second adjusts these values from the root to other nodes using precomputed information.
Steps:
- Perform a first DFS to calculate:
count[node]: number of nodes in the subtree rooted atnode.answer[0]: initial sum of distances from root node to all other nodes.- Execute a second DFS to propagate these sums effectively to the rest of the nodes and adjust them using the precomputed values:
- For each node transitioning from parent, calculate its sum based on the child's distance by subtracting its distance contributed and adding distances from non-subtree nodes.
Code:
- Time Complexity: O(N), as each edge and node is traversed a constant number of times.
- Space Complexity: O(N) for graph storage and arrays to store results and counts.