Shortest Path in a Weighted Tree

Shortest Path in a Weighted Tree - Problem

You are given an integer n and an undirected, weighted tree rooted at node 1 with n nodes numbered from 1 to n. This is represented by a 2D array edges of length n - 1, where edges[i] = [u_i, v_i, w_i] indicates an undirected edge from node u_i to v_i with weight w_i.

You are also given a 2D integer array queries of length q, where each queries[i] is either:

[1, u, v, w'] – Update the weight of the edge between nodes u and v to w', where (u, v) is guaranteed to be an edge present in edges.
[2, x] – Compute the shortest path distance from the root node 1 to node x.

Return an integer array answer, where answer[i] is the shortest path distance from node 1 to x for the i-th query of type [2, x].

Input & Output

Example 1 — Basic Tree Operations

$ Input: n = 4, edges = [[1,2,5],[1,3,3],[2,4,2]], queries = [[2,4],[1,1,3,10],[2,4]]

› Output: [7,7]

💡 Note: Initially: distance from 1 to 4 is 1→2→4 = 5+2 = 7. After updating edge (1,3) to weight 10, the distance from 1 to 4 remains 7 because the shortest path 1→2→4 doesn't use the updated edge (1,3).

Example 2 — Single Node Query

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

› Output: [8,3]

💡 Note: Initially: distance from 1 to 2 is 8. After updating edge (1,2) to weight 3, distance becomes 3.

Example 3 — Multiple Updates

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

› Output: [3,7,6]

💡 Note: Path 1→2→3: Initially 1+2=3, after first update 5+2=7, after second update 5+1=6.

Constraints

1 ≤ n ≤ 1000
edges.length = n - 1
1 ≤ u_i, v_i ≤ n
1 ≤ w_i ≤ 10⁶
1 ≤ queries.length ≤ 1000
For type 1 queries: 1 ≤ u, v ≤ n and 1 ≤ w' ≤ 10⁶
For type 2 queries: 1 ≤ x ≤ n

Visualization

Tap to expand

Asked in

G Google 45 M Microsoft 32 a Amazon 28

The key insight is that in a tree, there's exactly one path between any two nodes. We can precompute distances from the root and selectively update affected subtrees when edge weights change. The optimized approach uses caching with selective updates: Time O(n + q×k), Space O(n).

Common Approaches

✓ Brute Force DFS for Each Query

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

Build adjacency list, update edge weights as needed, then run DFS from root for each distance query. Simple but inefficient for many queries.

Precompute Distances with Updates

⏱️ Time: O(n + q × k) Space: O(n)

Precompute distances from root to all nodes initially. For update queries, only recalculate distances for nodes in the affected subtree using DFS.

Brute Force DFS for Each Query — Algorithm Steps

Build adjacency list from edges
For update queries, modify edge weights
For distance queries, run DFS from root to target

Visualization

Tap to expand

Step-by-Step Walkthrough

Build Graph

Create adjacency list from edges

Process Updates

Modify edge weights when needed

DFS Search

Traverse tree from root to find target distance

Code -

solution.c — C

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

#define MAX_N 1001

static int graph[MAX_N][MAX_N];
static int weights[MAX_N][MAX_N];
static int graphSize[MAX_N];

int dfs(int node, int target, int parent, int currentDist) {
    if (node == target) {
        return currentDist;
    }
    
    for (int i = 0; i < graphSize[node]; i++) {
        int neighbor = graph[node][i];
        if (neighbor != parent) {
            int weight = weights[node][i];
            int result = dfs(neighbor, target, node, currentDist + weight);
            if (result != -1) {
                return result;
            }
        }
    }
    return -1;
}

void updateEdgeWeight(int u, int v, int newWeight) {
    for (int i = 0; i < graphSize[u]; i++) {
        if (graph[u][i] == v) {
            weights[u][i] = newWeight;
            break;
        }
    }
    for (int i = 0; i < graphSize[v]; i++) {
        if (graph[v][i] == u) {
            weights[v][i] = newWeight;
            break;
        }
    }
}

int* solution(int n, int edges[][3], int edgesSize, int** queries, int* querySizes, int queriesSize, int* returnSize) {
    // Initialize graph
    for (int i = 1; i <= n; i++) {
        graphSize[i] = 0;
    }
    
    // Build graph
    for (int i = 0; i < edgesSize; i++) {
        int u = edges[i][0], v = edges[i][1], w = edges[i][2];
        graph[u][graphSize[u]] = v;
        weights[u][graphSize[u]] = w;
        graphSize[u]++;
        graph[v][graphSize[v]] = u;
        weights[v][graphSize[v]] = w;
        graphSize[v]++;
    }
    
    int* answer = malloc(queriesSize * sizeof(int));
    int answerIndex = 0;
    
    for (int i = 0; i < queriesSize; i++) {
        if (queries[i][0] == 1) {
            updateEdgeWeight(queries[i][1], queries[i][2], queries[i][3]);
        } else {
            int x = queries[i][1];
            if (x == 1) {
                answer[answerIndex++] = 0;
            } else {
                int dist = dfs(1, x, -1, 0);
                answer[answerIndex++] = dist;
            }
        }
    }
    
    *returnSize = answerIndex;
    return answer;
}

int parseNumber(const char* str, int* pos) {
    int num = 0;
    int sign = 1;
    if (str[*pos] == '-') {
        sign = -1;
        (*pos)++;
    }
    while (str[*pos] >= '0' && str[*pos] <= '9') {
        num = num * 10 + (str[*pos] - '0');
        (*pos)++;
    }
    return num * sign;
}

int main() {
    int n;
    scanf("%d", &n);
    getchar();
    
    char line[10000];
    fgets(line, sizeof(line), stdin);
    
    // Parse edges
    static int edges[1000][3];
    int edgesSize = 0;
    
    if (strcmp(line, "[]\n") != 0) {
        int pos = 1; // skip first '['
        while (pos < strlen(line) - 1) {
            if (line[pos] == '[') {
                pos++;
                edges[edgesSize][0] = parseNumber(line, &pos);
                pos++; // skip ','
                edges[edgesSize][1] = parseNumber(line, &pos);
                pos++; // skip ','
                edges[edgesSize][2] = parseNumber(line, &pos);
                edgesSize++;
                while (pos < strlen(line) && line[pos] != ']') pos++;
                pos++; // skip ']'
            } else {
                pos++;
            }
        }
    }
    
    fgets(line, sizeof(line), stdin);
    
    // Parse queries
    static int* queries[1000];
    static int querySizes[1000];
    static int queryData[1000][4];
    int queriesSize = 0;
    
    int pos = 1; // skip first '['
    while (pos < strlen(line) - 1) {
        if (line[pos] == '[') {
            pos++;
            queries[queriesSize] = queryData[queriesSize];
            
            int count = 0;
            while (pos < strlen(line) && line[pos] != ']') {
                if (line[pos] >= '0' && line[pos] <= '9') {
                    queryData[queriesSize][count] = parseNumber(line, &pos);
                    count++;
                } else {
                    pos++;
                }
            }
            querySizes[queriesSize] = count;
            queriesSize++;
            if (pos < strlen(line)) pos++; // skip ']'
        } else {
            pos++;
        }
    }
    
    int returnSize;
    int* result = solution(n, edges, edgesSize, queries, querySizes, queriesSize, &returnSize);
    
    printf("[");
    for (int i = 0; i < returnSize; i++) {
        if (i > 0) printf(",");
        printf("%d", result[i]);
    }
    printf("]\n");
    
    free(result);
    return 0;
}

Time & Space Complexity

Time Complexity

⏱️

O(q × n)

For each of q queries, potentially traverse all n nodes in DFS

✓ Linear Growth

Space Complexity

O(n)

Adjacency list stores n-1 edges, DFS recursion depth up to n

⚡ Linearithmic Space

23.4K Views

Medium Frequency

~35 min Avg. Time

892 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 for Each Query — Algorithm Steps

Visualization

Code -

Time & Space Complexity

Select Compiler