Find Weighted Median Node in Tree - Practice Coding Problems

Find Weighted Median Node in Tree - Problem

Imagine you're navigating through a network of interconnected cities, where each road has a specific travel cost. Your task is to find the weighted median node along any path between two cities.

You are given an integer n representing the number of nodes (0 to n-1) in an undirected, weighted tree rooted at node 0. The tree structure is defined by a 2D array edges where edges[i] = [u_i, v_i, w_i] represents a bidirectional edge between nodes u_i and v_i with weight w_i.

The weighted median node is the first node x on the path from node u to node v where the cumulative sum of edge weights from u to x is greater than or equal to half of the total path weight.

For each query [u_j, v_j], determine the weighted median node along the path from u_j to v_j. Return an array where each element corresponds to the weighted median node for the respective query.

Input & Output

example_1.py — Basic Tree Query

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

› Output: [1]

💡 Note: For the path from node 3 to node 4: 3 → 1 → 4. The edge weights are [3, 5] with total weight 8. The cumulative sum reaches 3 at node 1, and 3 + 5 = 8 at node 4. Since 8/2 = 4, and cumulative weight 8 ≥ 4, the weighted median is node 4. However, since we want the first node where cumulative ≥ half, we need to check: at node 1, cumulative is 3 < 4, but at the next edge (1→4), cumulative becomes 8 ≥ 4, so the weighted median is node 4.

example_2.py — Multiple Queries

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

› Output: [2,2]

💡 Note: For query [0,3]: Path 0→1→2→3, weights [2,3,1], total=6, target=3. Cumulative: 2 at node 1, 5 at node 2 (≥3), so median is node 2. For query [1,3]: Path 1→2→3, weights [3,1], total=4, target=2. Cumulative: 3 at node 2 (≥2), so median is node 2.

example_3.py — Single Node Path

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

› Output: [1]

💡 Note: When source and destination are the same node, the weighted median is the node itself.

Constraints

2 ≤ n ≤ 10⁴
edges.length == n - 1
0 ≤ u_i, v_i < n
1 ≤ w_i ≤ 10⁶
1 ≤ queries.length ≤ 10⁴
The graph forms a valid tree (connected and acyclic)

Visualization

Tap to expand

Understanding the Visualization

Build Tree Structure

Create adjacency list representation of the weighted tree

Find Path Between Nodes

Use tree traversal or LCA to find the unique path between query nodes

Calculate Total Weight

Sum all edge weights along the path to get total cost

Find Weighted Median

Traverse path accumulating weights until reaching ≥ half of total weight

Key Takeaway

🎯 Key Insight: In a tree, there's exactly one path between any two nodes, so we can efficiently find weighted medians using LCA preprocessing and path reconstruction.

Asked in

G Google 38 a Amazon 31 f Meta 24 ⊞ Microsoft 19

The optimal approach uses LCA with Binary Lifting preprocessing to handle multiple queries efficiently. We preprocess the tree in O(N log N) time, then answer each query in O(log N) time by finding the LCA and reconstructing the path to calculate the weighted median. This is much more efficient than the brute force approach which takes O(Q × N) time for Q queries.

Common Approaches

Approach	Time	Space	Notes
✓ LCA + Binary Lifting (Optimal)	O(N log N + Q log N)	O(N log N)	Use LCA with binary lifting for efficient path queries and weighted median calculation
Brute Force Path Finding	O(Q × N)	O(N)	For each query, find the path using DFS and calculate weighted median by iterating through the path

LCA + Binary Lifting (Optimal) — Algorithm Steps

Preprocess the tree with binary lifting for fast LCA queries
For each query, find LCA of the two nodes
Reconstruct path from u to LCA and LCA to v
Calculate cumulative weights and find weighted median efficiently

Visualization

Tap to expand

Step-by-Step Walkthrough

Preprocess Tree

Build binary lifting table with parent pointers

Find LCA

Use binary lifting to find lowest common ancestor

Reconstruct Path

Build path from u to LCA to v

Calculate Median

Find weighted median along the reconstructed path

Code -

solution.c — C

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

#define MAX_N 1000
#define MAX_LOG 20

typedef struct {
    int node, weight;
} Edge;

typedef struct {
    Edge edges[MAX_N];
    int count;
} AdjList;

AdjList graph[MAX_N];
int parent[MAX_N][MAX_LOG];
int depth[MAX_N];
int LOG;

void dfs(int node, int par, int d) {
    parent[node][0] = par;
    depth[node] = d;
    
    for (int i = 1; i < LOG; i++) {
        if (parent[node][i-1] != -1) {
            parent[node][i] = parent[parent[node][i-1]][i-1];
        }
    }
    
    for (int i = 0; i < graph[node].count; i++) {
        if (graph[node].edges[i].node != par) {
            dfs(graph[node].edges[i].node, node, d + 1);
        }
    }
}

int lca(int u, int v) {
    if (depth[u] < depth[v]) {
        int temp = u; u = v; v = temp;
    }
    
    int diff = depth[u] - depth[v];
    for (int i = 0; i < LOG; i++) {
        if ((diff >> i) & 1) {
            u = parent[u][i];
        }
    }
    
    if (u == v) return u;
    
    for (int i = LOG - 1; i >= 0; i--) {
        if (parent[u][i] != parent[v][i]) {
            u = parent[u][i];
            v = parent[v][i];
        }
    }
    
    return parent[u][0];
}

int findWeightedMedian(int u, int v) {
    if (u == v) return u;
    
    int lcaNode = lca(u, v);
    int path[MAX_N], weights[MAX_N];
    int pathLen = 0, weightLen = 0;
    
    int curr = u;
    while (curr != lcaNode) {
        path[pathLen++] = curr;
        for (int i = 0; i < graph[curr].count; i++) {
            if (depth[graph[curr].edges[i].node] < depth[curr] && 
                graph[curr].edges[i].node == parent[curr][0]) {
                weights[weightLen++] = graph[curr].edges[i].weight;
                break;
            }
        }
        curr = parent[curr][0];
    }
    
    path[pathLen++] = lcaNode;
    
    int tempPath[MAX_N], tempWeights[MAX_N];
    int tempPathLen = 0, tempWeightLen = 0;
    
    curr = v;
    while (curr != lcaNode) {
        tempPath[tempPathLen++] = curr;
        for (int i = 0; i < graph[curr].count; i++) {
            if (depth[graph[curr].edges[i].node] < depth[curr] && 
                graph[curr].edges[i].node == parent[curr][0]) {
                tempWeights[tempWeightLen++] = graph[curr].edges[i].weight;
                break;
            }
        }
        curr = parent[curr][0];
    }
    
    for (int i = tempPathLen - 1; i >= 0; i--) {
        path[pathLen++] = tempPath[i];
    }
    for (int i = tempWeightLen - 1; i >= 0; i--) {
        weights[weightLen++] = tempWeights[i];
    }
    
    int totalWeight = 0;
    for (int i = 0; i < weightLen; i++) {
        totalWeight += weights[i];
    }
    
    double target = totalWeight / 2.0;
    int cumulative = 0;
    
    for (int i = 0; i < weightLen; i++) {
        cumulative += weights[i];
        if (cumulative >= target) {
            return path[i + 1];
        }
    }
    
    return path[pathLen - 1];
}

int* solution(int n, int** edges, int edgesSize, int** queries, int queriesSize, int* returnSize) {
    LOG = (int)ceil(log2(n)) + 1;
    
    for (int i = 0; i < n; i++) {
        graph[i].count = 0;
        for (int j = 0; j < LOG; j++) {
            parent[i][j] = -1;
        }
    }
    
    for (int i = 0; i < edgesSize; i++) {
        int u = edges[i][0], v = edges[i][1], w = edges[i][2];
        graph[u].edges[graph[u].count++] = (Edge){v, w};
        graph[v].edges[graph[v].count++] = (Edge){u, w};
    }
    
    dfs(0, -1, 0);
    
    int* result = (int*)malloc(queriesSize * sizeof(int));
    *returnSize = queriesSize;
    
    for (int i = 0; i < queriesSize; i++) {
        result[i] = findWeightedMedian(queries[i][0], queries[i][1]);
    }
    
    return result;
}

Time & Space Complexity

Time Complexity

⏱️

O(N log N + Q log N)

O(N log N) for preprocessing with binary lifting, O(log N) per query for LCA

⚡ Linearithmic

Space Complexity

O(N log N)

Space for binary lifting table and tree preprocessing

✓ Linear Space

43.7K Views

Medium Frequency

~25 min Avg. Time

1.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

LCA + Binary Lifting (Optimal) — Algorithm Steps

Visualization

Code -

Time & Space Complexity

Select Compiler