Code Implementation

class Solution:
    def sumOfDistancesInTree(self, N, edges):
        from collections import defaultdict
        tree = defaultdict(list)
        for u, v in edges:
            tree[u].append(v)
            tree[v].append(u)
        
        res = [0] * N
        count = [1] * N  # Each node counts as itself

        def post_order(node, parent):
            for child in tree[node]:
                if child != parent:
                    post_order(child, node)
                    count[node] += count[child]
                    res[node] += res[child] + count[child]

        def pre_order(node, parent):
            for child in tree[node]:
                if child != parent:
                    res[child] = res[node] - count[child] + (N - count[child])
                    pre_order(child, node)

        post_order(0, -1)
        pre_order(0, -1)
        return res
      

class Solution {
public:
    vector<int> sumOfDistancesInTree(int N, vector<vector<int>>& edges) {
        vector<vector<int>> tree(N);
        for (const auto& e : edges) {
            tree[e[0]].push_back(e[1]);
            tree[e[1]].push_back(e[0]);
        }
        vector<int> res(N, 0), count(N, 1);

        function<void(int, int)> post_order = [&](int node, int parent) {
            for (int child : tree[node]) {
                if (child == parent) continue;
                post_order(child, node);
                count[node] += count[child];
                res[node] += res[child] + count[child];
            }
        };

        function<void(int, int)> pre_order = [&](int node, int parent) {
            for (int child : tree[node]) {
                if (child == parent) continue;
                res[child] = res[node] - count[child] + (N - count[child]);
                pre_order(child, node);
            }
        };

        post_order(0, -1);
        pre_order(0, -1);
        return res;
    }
};
      

class Solution {
    List<List<Integer>> tree;
    int[] res, count;
    public int[] sumOfDistancesInTree(int N, int[][] edges) {
        tree = new ArrayList<>();
        for (int i = 0; i < N; ++i) tree.add(new ArrayList<>());
        for (int[] e : edges) {
            tree.get(e[0]).add(e[1]);
            tree.get(e[1]).add(e[0]);
        }
        res = new int[N];
        count = new int[N];
        Arrays.fill(count, 1);

        postOrder(0, -1);
        preOrder(0, -1, N);
        return res;
    }
    private void postOrder(int node, int parent) {
        for (int child : tree.get(node)) {
            if (child == parent) continue;
            postOrder(child, node);
            count[node] += count[child];
            res[node] += res[child] + count[child];
        }
    }
    private void preOrder(int node, int parent, int N) {
        for (int child : tree.get(node)) {
            if (child == parent) continue;
            res[child] = res[node] - count[child] + (N - count[child]);
            preOrder(child, node, N);
        }
    }
}
      

var sumOfDistancesInTree = function(N, edges) {
    const tree = Array.from({length: N}, () => []);
    for (const [u, v] of edges) {
        tree[u].push(v);
        tree[v].push(u);
    }
    const res = Array(N).fill(0);
    const count = Array(N).fill(1);

    function postOrder(node, parent) {
        for (const child of tree[node]) {
            if (child === parent) continue;
            postOrder(child, node);
            count[node] += count[child];
            res[node] += res[child] + count[child];
        }
    }

    function preOrder(node, parent) {
        for (const child of tree[node]) {
            if (child === parent) continue;
            res[child] = res[node] - count[child] + (N - count[child]);
            preOrder(child, node);
        }
    }

    postOrder(0, -1);
    preOrder(0, -1);
    return res;
};
      

Problem Description

Given a tree with N nodes labeled from 0 to N-1 and an array edges where each element [u, v] represents an undirected edge between nodes u and v, your task is to return an array answer of length N such that answer[i] is the sum of the distances between node i and all other nodes in the tree.

• The input graph is a tree, i.e., it is connected and has no cycles.
• There is exactly one valid solution for each input.
• Each node must be considered as the root once, and the sum of distances to all other nodes is calculated for that root.
• Constraints:
  • 1 <= N <= 3 * 10^4
  • edges.length == N - 1
  • 0 <= u, v < N

Thought Process

At first glance, the problem seems to require calculating the sum of distances from every node to every other node. A naive approach would be to perform a breadth-first search (BFS) or depth-first search (DFS) from each node, summing up distances to all others. However, this would result in a time complexity of O(N^2), which is not feasible for large values of N (up to 30,000).

To optimize, we need to find a way to reuse computations between nodes. Since a tree is a connected acyclic graph, we can exploit its structure to compute the answer for all nodes efficiently by performing two passes:

• A post-order traversal to gather subtree information (number of descendants and partial distances).
• A pre-order traversal to propagate and adjust the distance sums from parent to child nodes.

Solution Approach

We use a two-pass DFS approach to efficiently compute the sum of distances for every node:

1. Build the Tree:
   • Represent the tree as an adjacency list for efficient traversal.
2. First Pass - Post-order DFS:
   • Start from an arbitrary root (e.g., node 0).
   • For each node, compute:
     • count[node]: the number of nodes in the subtree rooted at node (including itself).
     • res[node]: the sum of distances from node to all nodes in its subtree.
   • For each child, after the DFS, update:
     • count[node] += count[child]
     • res[node] += res[child] + count[child] (distance to each node in the child subtree increases by 1 from the parent)
3. Second Pass - Pre-order DFS:
   • Propagate the results from parent to children.
   • For each child, compute:
     • res[child] = res[parent] - count[child] + (N - count[child])
   • This formula shifts the root from parent to child:
     • Distances to nodes in the child's subtree decrease by 1 (closer to new root).
     • Distances to nodes not in the child's subtree increase by 1 (farther from new root).
   • Repeat recursively for all children.
4. Return the result array res after both traversals.

This approach leverages the tree's recursive properties and computes the answer in O(N) time.

Example Walkthrough

Consider N = 6 and edges = [[0,1],[0,2],[2,3],[2,4],[2,5]]. The tree structure is:

        0
       / \
      1   2
         /|\
        3 4 5
  

1. Post-order traversal:
   • Start at node 0, traverse down to leaves.
   • For leaf nodes (1, 3, 4, 5), count = 1 and res = 0.
   • For node 2:
     • Children: 3, 4, 5
     • count[2] = 1 + 1 + 1 + 1 = 4
     • res[2] = 0 + 1 + 0 + 1 + 0 + 1 = 3
   • For node 0:
     • Children: 1, 2
     • count[0] = 1 + 1 + 4 = 6
     • res[0] = 0 + 1 + 3 + 4 = 8
2. Pre-order traversal:
   • Propagate res from parent to children.
   • For node 1:
     • res[1] = res[0] - count[1] + (N - count[1]) = 8 - 1 + 5 = 12
   • For node 2:
     • res[2] = res[0] - count[2] + (N - count[2]) = 8 - 4 + 2 = 6
   • For node 3:
     • res[3] = res[2] - count[3] + (N - count[3]) = 6 - 1 + 5 = 10
   • Similarly for nodes 4 and 5.
3. Final result: [8,12,6,10,10,10]

Time and Space Complexity

• Brute-force approach:
  • For each node, run BFS/DFS to compute all distances: O(N^2) time.
• Optimized approach (two-pass DFS):
  • Each node and edge is visited a constant number of times: O(N) time.
  • Space for adjacency list and helper arrays: O(N) space.

The optimized approach is efficient and scalable even for large trees.

Summary

This problem can be solved efficiently by leveraging the recursive structure of trees. By performing a post-order traversal, we gather subtree sizes and initial distances. Then, a pre-order traversal allows us to adjust these distances for every possible root in a single pass. The key insight is that the sum of distances for a child can be derived from its parent using subtree counts, resulting in a clean and elegant O(N) solution.