Sum of Distances in Tree - Practice Coding Problems

Sum of Distances in Tree - Problem

You are given an undirected connected tree with n nodes labeled from 0 to n - 1 and exactly n - 1 edges. Your task is to calculate the sum of distances from each node to all other nodes in the tree.

Given:

An integer n representing the number of nodes
An array edges where edges[i] = [a_i, b_i] indicates an edge between nodes a_i and b_i

Goal: Return an array answer of length n where answer[i] is the sum of distances between the i-th node and all other nodes in the tree.

Note: The distance between two nodes is the number of edges in the shortest path between them.

Input & Output

example_1.py — Basic Tree

$ Input: n = 6, edges = [[0,1],[0,2],[2,3],[2,4],[2,5]]

› Output: [8,12,6,10,10,10]

💡 Note: Tree looks like: 1-0-2-3,4,5. From node 0: distances are [0,1,1,2,2,2] so sum=8. From node 2: distances are [0,1,2,1,1,1] so sum=6.

example_2.py — Linear Tree

$ Input: n = 3, edges = [[1,0],[2,1]]

› Output: [3,2,3]

💡 Note: Linear tree: 0-1-2. From node 0: distances [0,1,2] sum=3. From node 1: distances [1,0,1] sum=2. From node 2: distances [2,1,0] sum=3.

example_3.py — Single Node

$ Input: n = 1, edges = []

› Output: [0]

💡 Note: Only one node, so distance sum is 0.

Visualization

Tap to expand

Understanding the Visualization

Build the Tree

Create adjacency list representation of the tree from given edges

First DFS Pass

Calculate subtree sizes and sum of distances from an arbitrary root (node 0)

Second DFS Pass

Use rerooting: when moving from parent to child, some nodes get closer, others get farther

Mathematical Transform

Apply the formula to efficiently calculate answer for each node

Key Takeaway

🎯 Key Insight: Instead of recalculating distances from scratch for each node, we can mathematically transform the answer from parent to child using subtree size information!

Time & Space Complexity

Time Complexity

⏱️

O(n)

Two DFS passes, each visiting every node exactly once

✓ Linear Growth

Space Complexity

O(n)

Space for adjacency list, count array, and recursion stack

⚡ Linearithmic Space

Constraints

1 ≤ n ≤ 3 × 10⁴
edges.length == n - 1
edges[i].length == 2
0 ≤ a_i, b_i < n
a_i ≠ b_i
The given input represents a valid tree

Asked in

G Google 42 a Amazon 28 f Meta 35 ⊞ Microsoft 18

The optimal solution uses two-pass dynamic programming with a rerooting technique. First pass calculates subtree sizes and distances from root, second pass efficiently redistributes this information to all nodes using the mathematical relationship between parent and child answers.

Common Approaches

Approach	Time	Space	Notes
✓ Two-Pass Dynamic Programming (Optimal)	O(n)	O(n)	Use tree DP with rerooting technique to solve in O(n) time
Brute Force (DFS from Each Node)	O(n²)	O(n)	Run DFS from each node to calculate distances to all other nodes

Two-Pass Dynamic Programming (Optimal) — Algorithm Steps

Build adjacency list from edges
First DFS: Calculate count[i] (subtree size) and answer[0] (sum of distances from root)
Second DFS: Use rerooting technique to calculate answer[i] for all other nodes
When moving from parent to child, adjust the answer using the formula
Return the final answer array

Visualization

Tap to expand

Step-by-Step Walkthrough

First DFS Pass

Calculate subtree sizes and sum of distances from root

Second DFS Pass

Use rerooting technique to calculate answers for all nodes

Mathematical Transform

Apply formula: answer[child] = answer[parent] - count[child] + (n - count[child])

Code -

solution.c — C

#include <stdio.h>
#include <stdlib.h>

int** graph;
int* graphSize;
int* count;
int* answer;
int n;

void dfs1(int node, int parent) {
    for (int i = 0; i < graphSize[node]; i++) {
        int child = graph[node][i];
        if (child != parent) {
            dfs1(child, node);
            count[node] += count[child];
            answer[node] += answer[child] + count[child];
        }
    }
}

void dfs2(int node, int parent) {
    for (int i = 0; i < graphSize[node]; i++) {
        int child = graph[node][i];
        if (child != parent) {
            answer[child] = answer[node] - count[child] + (n - count[child]);
            dfs2(child, node);
        }
    }
}

int* sumOfDistancesInTree(int n_param, int** edges, int edgesSize, int* edgesColSize, int* returnSize) {
    n = n_param;
    *returnSize = n;
    
    graph = (int**)malloc(n * sizeof(int*));
    graphSize = (int*)calloc(n, sizeof(int));
    count = (int*)malloc(n * sizeof(int));
    answer = (int*)calloc(n, sizeof(int));
    
    for (int i = 0; i < n; i++) {
        graph[i] = (int*)malloc(n * sizeof(int));
        count[i] = 1;
    }
    
    for (int i = 0; i < edgesSize; i++) {
        int a = edges[i][0], b = edges[i][1];
        graph[a][graphSize[a]++] = b;
        graph[b][graphSize[b]++] = a;
    }
    
    dfs1(0, -1);
    dfs2(0, -1);
    
    return answer;
}

int main() {
    // Parse input and call solution
    return 0;
}

Time & Space Complexity

Time Complexity

⏱️

O(n)

Two DFS passes, each visiting every node exactly once

✓ Linear Growth

Space Complexity

O(n)

Space for adjacency list, count array, and recursion stack

⚡ Linearithmic Space

Constraints

1 ≤ n ≤ 3 × 10⁴
edges.length == n - 1
edges[i].length == 2
0 ≤ a_i, b_i < n
a_i ≠ b_i
The given input represents a valid tree

89.7K Views

Medium-High Frequency

~25 min Avg. Time

2.8K Likes

Ln 1, Col 1

Smart Actions

💡 Explanation

AI Ready

💡 Suggestion Tab to accept Esc to dismiss

// Output will appear here after running code

Code Editor Closed

Click the red button to reopen

Input & Output

Visualization

Time & Space Complexity

Related Problems

Constraints

Common Approaches

Two-Pass Dynamic Programming (Optimal) — Algorithm Steps

Visualization

Code -

Time & Space Complexity

Constraints

Select Compiler