Sum Of Distances - Leetcode Solution

The "Sum of Distances in Tree" problem on LeetCode asks us to find the sum of distances between all pairs of nodes in a tree. In other words, for each node in the tree, we need to calculate the sum of distances to all other nodes in the tree.

A straightforward approach to solve this problem is to perform a depth-first traversal of the tree starting from each node. For each node, we can keep track of the distance to all other nodes using a separate array. Once we have computed the distance arrays for all nodes, we can sum up the distance values across all nodes to get the final solution.

Let's discuss the detailed solution.

1. Create an adjacency list to represent the tree.

2. Initialize two arrays: "count" and "dist". "count[i]" will store the count of nodes in the subtree rooted at node "i", while "dist[i]" will store the sum of distances from node "i" to all other nodes in the subtree rooted at node "i".

3. Perform a depth-first traversal of the tree starting from node 0. For every child "j" of node "i", call the DFS function recursively on "j" and update "count[i]" and "dist[i]".

   a. After we have finished exploring the subtree rooted at "j", we can update "count[i]" and "dist[i]" based on the values returned from the function.

   b. "count[i]" will be incremented by "count[j]" (since all nodes in the subtree rooted at "j" are also in the subtree rooted at "i").

   c. "dist[i]" will be incremented by "dist[j] + count[j]" (since we need to add the distance from "i" to "j" for all nodes in the subtree rooted at "j").

4. Once we have computed the "count" and "dist" arrays for all nodes, we can calculate the final solution.

   a. Initialize an array "res" to store the final result. "res[i]" will store the sum of distances from node "i" to all other nodes in the tree.

   b. Perform another depth-first traversal starting from node 0. For every child "j" of node "i", call the DFS function recursively on "j" and update "res[i]".

   i. Since all nodes in the subtree rooted at "i" are also in the subtree rooted at "j", we can subtract "count[j]" from "res[i]".

   ii. Similarly, since all nodes in the subtree rooted at "j" are also in the subtree rooted at "i", we can add "count[0] - count[j]" to "res[i]".

   iii. Finally, we need to add "dist[j]" to "res[i]" (since all nodes in the subtree rooted at "j" are at a distance of 1 from "j").

   c. Once we have calculated the "res" array for all nodes, we can return it as the final solution.

Here is the Python code for the solution:

class Solution:
    def sumOfDistancesInTree(self, N: int, edges: List[List[int]]) -> List[int]:
        tree = collections.defaultdict(set)
        for i, j in edges:
            tree[i].add(j)
            tree[j].add(i)
        
        count = [1] * N
        dist = [0] * N
        
        def dfs1(node, parent):
            for child in tree[node]:
                if child != parent:
                    dfs1(child, node)
                    count[node] += count[child]
                    dist[node] += dist[child] + count[child]
            
        dfs1(0, -1)
        
        res = [0] * N
        
        def dfs2(node, parent):
            res[node] = dist[node]
            for child in tree[node]:
                if child != parent:
                    res[node] -= count[child]
                    res[node] += N - count[child]
                    res[node] += dist[child]
                    dfs2(child, node)
            
        dfs2(0, -1)
        return res

The time complexity of this solution is O(N), where N is the number of nodes in the tree. This is because we perform two depth-first traversals of the tree, each of which visits every node exactly once. The space complexity is also O(N), since we store the "count" and "dist" arrays, as well as the recursive call stack for the depth-first traversals.

1