Sum of Distances in Tree

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: [5,3,5,5]

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

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

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

✓ Brute Force DFS

⏱️ Time: O(n²) Space: O(n)

For each node, perform a depth-first search to calculate the distance to all other nodes. This gives us the exact answer but is inefficient for large trees.

Tree DP with Rerooting

⏱️ Time: O(n) Space: O(n)

First DFS calculates subtree sizes and distances when node 0 is root. Second DFS uses rerooting technique to efficiently calculate distances for all other nodes as roots.

Brute Force DFS — Algorithm Steps

For each node i from 0 to n-1
Run DFS from node i to calculate distances to all other nodes
Sum up all distances and store in result[i]

Visualization

Tap to expand

Step-by-Step Walkthrough

Start from Node 0

Run DFS to calculate distances to all other nodes

Repeat for All Nodes

Do the same process for each node as root

Collect Results

Store sum of distances for each starting node

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;

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 dfs(AdjList* graph, int node, int dist, bool* visited, int* result, int root, int n) {
    visited[node] = true;
    result[root] += dist;
    for (int i = 0; i < graph[node].size; i++) {
        int neighbor = graph[node].nodes[i];
        if (!visited[neighbor]) {
            dfs(graph, neighbor, dist + 1, visited, result, root, n);
        }
    }
}

int* solution(int n, int edges[][2], int edgesSize, int* returnSize) {
    *returnSize = n;
    int* result = calloc(n, sizeof(int));
    
    // Build adjacency list
    AdjList* 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]);
    }
    
    // For each node, calculate sum of distances to all other nodes
    for (int root = 0; root < n; root++) {
        bool* visited = calloc(n, sizeof(bool));
        dfs(graph, root, 0, visited, result, root, n);
        free(visited);
    }
    
    // Cleanup
    for (int i = 0; i < n; i++) {
        free(graph[i].nodes);
    }
    free(graph);
    
    return result;
}

int main() {
    int n;
    scanf("%d", &n);
    
    char line[10000];
    scanf(" %[^\n]", line);
    
    // Simple parsing for edges array
    int edges[1000][2];
    int edgesSize = 0;
    
    char* ptr = line + 1; // Skip first '['
    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²)

For each of n nodes, we perform DFS visiting all n nodes

✓ Linear Growth

Space Complexity

O(n)

Recursion stack depth up to n for DFS traversal

⚡ 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

Brute Force DFS — Algorithm Steps

Visualization

Code -

Time & Space Complexity

Select Compiler