Problem statement

https://leetcode.com/problems/sum-of-distances-in-tree/

Solution

Let us first create graph G from our edges. Let us also use two different dfs functions:

  1. dfs(node, parent, depth), where node is current node we are in, parent is parent of current node and depth is depth of node. This function is used to traverse tree and for each node find two values: depth of node (how far it is from root) and also weight of node: how many nodes in subtree with given node. For example, for tree [[0,1],[0,2],[2,3],[2,4],[2,5]] weights will be [6, 1, 4, 1, 1, 1].
  2. dfs2(node, parent, w), where node and parent current node and its parent and w in answer for current node. How we can find answer for neib if we already now answer for node? It is w + N - 2*weights[neib], because when we moved from node to neib not a lot changed: for weights[neib] number of nodes distanced increased by 1 and for N - weights[neib] number of nodes distances decreased by 1.

Finally, we just run dfs2(0, -1, sum(depths)).

Complexity

Time and space complexity is O(n).

Code

class Solution:
    def sumOfDistancesInTree(self, N, edges):
        G = defaultdict(set)
        for i, j in edges:
            G[i].add(j)
            G[j].add(i)
        
        def dfs(node, parent, depth):
            ans = 1
            for neib in G[node]:
                if neib != parent:
                    ans += dfs(neib, node, depth + 1)
            weights[node] = ans
            depths[node] = depth
            return ans
        
        def dfs2(node, parent, w):
            ans[node] = w
            for neib in G[node]:
                if neib != parent:
                    dfs2(neib, node, w + N - 2*weights[neib])
        
        weights, depths, ans = [0]*N, [0]*N, [0]*N
        dfs(0, -1, 0)
        dfs2(0, -1, sum(depths))
        
        return ans