Sum of Distances in Tree - Practice Coding Problems

Sum of Distances in Tree - Problem

There is an undirected connected tree with n nodes labeled from 0 to n - 1 and n - 1 edges.

You are given the integer n and the array edges where edges[i] = [a_i, b_i] indicates that there is an edge between nodes a_i and b_i in the tree.

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

Input & Output

Example 1 — Simple Tree

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

› Output: [6,4,4,4]

💡 Note: Tree looks like: 0-1-2 with 3 also connected to 1. From node 0: distances are [0,1,2,2] sum=5. From node 1: distances are [1,0,1,1] sum=3. Continue for all nodes.

Example 2 — Linear Tree

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

› Output: [3,2,3]

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

Example 3 — Single Edge

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

› Output: [1,1]

💡 Note: Two nodes connected: 0-1. From each node, distance to other is 1.

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.

Visualization

Tap to expand

Understanding the Visualization

Input Tree

Undirected tree with n nodes and n-1 edges

Calculate Distances

For each node, find sum of distances to all other nodes

Output Array

Array where result[i] = sum of distances from node i

Key Takeaway

🎯 Key Insight: Use rerooting technique to avoid recalculating distances from scratch for each node

Asked in

G Google 15 F Facebook 12 A Amazon 8

The key insight is using tree dynamic programming with a rerooting technique. First DFS calculates subtree information, second DFS efficiently computes distances for all nodes. Optimal approach uses two-pass DFS with O(n) time and O(n) space.

Common Approaches

Approach	Time	Space	Notes
✓ Tree DP with Rerooting	O(n)	O(n)	Use dynamic programming to calculate distances efficiently in two DFS passes
Brute Force DFS	O(n²)	O(n)	Run DFS from each node to calculate sum of distances

Tree DP with Rerooting — Algorithm Steps

Build adjacency list from edges
First DFS: calculate subtree sizes and answer for node 0
Second DFS: use rerooting to calculate answers for all other nodes

Visualization

Tap to expand

Step-by-Step Walkthrough

First DFS

Calculate subtree sizes and distances for root 0

Rerooting

Use formula to adjust distances when changing root

Final Result

All distances calculated in linear time

Code -

solution.c — C

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

typedef struct {
    int* nodes;
    int size;
    int capacity;
} AdjList;

static AdjList* graph;
static int* count;
static int* ans;
static int n;

void addEdge(AdjList* graph, int u, int v) {
    if (graph[u].size == graph[u].capacity) {
        graph[u].capacity *= 2;
        graph[u].nodes = realloc(graph[u].nodes, graph[u].capacity * sizeof(int));
    }
    graph[u].nodes[graph[u].size++] = v;
}

void dfs1(int node, int parent) {
    count[node] = 1;
    for (int i = 0; i < graph[node].size; i++) {
        int child = graph[node].nodes[i];
        if (child != parent) {
            dfs1(child, node);
            count[node] += count[child];
            ans[node] += ans[child] + count[child];
        }
    }
}

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

int* solution(int n, int edges[][2], int edgesSize, int* returnSize) {
    *returnSize = n;
    ::n = n;
    
    int* result = calloc(n, sizeof(int));
    count = calloc(n, sizeof(int));
    ans = calloc(n, sizeof(int));
    
    // Build adjacency list
    graph = malloc(n * sizeof(AdjList));
    for (int i = 0; i < n; i++) {
        graph[i].nodes = malloc(10 * sizeof(int));
        graph[i].size = 0;
        graph[i].capacity = 10;
    }
    
    for (int i = 0; i < edgesSize; i++) {
        addEdge(graph, edges[i][0], edges[i][1]);
        addEdge(graph, edges[i][1], edges[i][0]);
    }
    
    dfs1(0, -1);
    dfs2(0, -1);
    
    for (int i = 0; i < n; i++) {
        result[i] = ans[i];
    }
    
    // Cleanup
    for (int i = 0; i < n; i++) {
        free(graph[i].nodes);
    }
    free(graph);
    free(count);
    free(ans);
    
    return result;
}

int main() {
    int n;
    scanf("%d", &n);
    
    char line[10000];
    scanf(" %[^\n]", line);
    
    int edges[1000][2];
    int edgesSize = 0;
    
    char* ptr = line + 1;
    while (*ptr && *ptr != ']') {
        if (*ptr == '[') {
            ptr++;
            edges[edgesSize][0] = atoi(ptr);
            while (*ptr != ',') ptr++;
            ptr++;
            edges[edgesSize][1] = atoi(ptr);
            edgesSize++;
            while (*ptr != ']') ptr++;
            ptr++;
        } else {
            ptr++;
        }
    }
    
    int returnSize;
    int* result = solution(n, edges, edgesSize, &returnSize);
    
    printf("[");
    for (int i = 0; i < returnSize; i++) {
        printf("%d", result[i]);
        if (i < returnSize - 1) printf(",");
    }
    printf("]\n");
    
    free(result);
    return 0;
}

Time & Space Complexity

Time Complexity

⏱️

O(n)

Two DFS traversals, each visiting every node exactly once

✓ Linear Growth

Space Complexity

O(n)

Adjacency list storage and recursion stack depth

⚡ Linearithmic Space

89.5K Views

Medium Frequency

~35 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

Constraints

Visualization

Related Problems

Common Approaches

Tree DP with Rerooting — Algorithm Steps

Visualization

Code -

Time & Space Complexity

Select Compiler