Sum Of Distances In Tree - Leetcode Solution

The Sum of Distances in Tree problem on LeetCode asks you to find the sum of distances between all pairs of nodes in a tree. The input is an array of edges that represent the tree. The output is an array of integers that represent the sum of distances for each node in the tree.

To solve this problem, we can use two traversals of the tree: one to calculate the subtree size and the distance from the root for each node, and another to calculate the sum of distances for each node.

Algorithm:

1. Create an adjacency list that represents the tree from the given edges.
2. Initialize two arrays:

• subtree_size: An array that stores the size of the subtree rooted at each node.
• distance: An array that stores the distance of each node from the root.

3. Initialize the distance array with 0 for the root node.
4. Calculate the subtree_size and distance for all nodes using a depth-first search (DFS) traversal: a. Start the DFS from the root node. b. For each child of a node u, add 1 to subtree_size[u] and calculate distance[v] = distance[u] + 1. c. Recursively traverse all the children. d. After traversing all the children, add subtree_size[v] to subtree_size[u].
5. Initialize another array ans with 0.
6. Traverse the tree again using DFS to calculate the ans array: a. Start the DFS from the root node. b. For each child of a node u, calculate ans[v] = ans[u] + (n - subtree_size[v]) - subtree_size[v]. c. Recursively traverse all the children.
7. Return the ans array.

The time complexity of this algorithm is O(N), where N is the number of nodes in the tree. This is because we traverse the tree twice using a DFS traversal, and at each node, we perform constant time operations.

Here is the implementation of the solution in Python:

class Solution:
    def sumOfDistancesInTree(self, N: int, edges: List[List[int]]) -> List[int]:
        # Step 1
        adj = [[] for _ in range(N)]
        for u, v in edges:
            adj[u].append(v)
            adj[v].append(u)
        
        # Step 2
        subtree_size = [1] * N
        distance = [0] * N
        
        # Step 3
        def dfs1(u, parent):
            for v in adj[u]:
                if v != parent:
                    distance[v] = distance[u] + 1
                    dfs1(v, u)
                    subtree_size[u] += subtree_size[v]
        
        dfs1(0, -1)
        
        # Step 4
        ans = [0] * N
        
        # Step 6
        def dfs2(u, parent):
            for v in adj[u]:
                if v != parent:
                    ans[v] = ans[u] + (N - subtree_size[v]) - subtree_size[v]
                    dfs2(v, u)
        
        dfs2(0, -1)
        
        # Step 7
        ans[0] = sum(distance)
        for i in range(1, N):
            ans[i] += ans[0] - 2 * distance[i]
        
        return ans

In this implementation, we used a recursive DFS traversal instead of an iterative DFS traversal for simplicity.

1