Find Weighted Median Node in Tree

Hard

#3192Time: O(Q * N), where N is the number of nodes and Q is the number of queries. For each query, the traversal (BFS/DFS) to find the path can take up to O(N + E) = O(N) time in the worst case for a tree. Reconstructing and traversing the path also takes up to O(N) time.Space: O(N) to store the adjacency list, the visited array, the parent map, and the path for each query.

Patterns

Dynamic Programming

Algorithms

Binary Search Depth-First Search

Data structures

Array Tree

Prompt

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

The weighted median node is defined as the first node x on the path from u_i to v_i such that the sum of edge weights from u_i to x is greater than or equal to half of the total path weight.

You are given a 2D integer array queries. For each queries[j] = [u_j, v_j], determine the weighted median node along the path from u_j to v_j.

Return an array ans, where ans[j] is the node index of the weighted median for queries[j].

Example 1:

Input: n = 2, edges = [[0,1,7]], queries = [[1,0],[0,1]]

Output: [0,1]

Explanation:

Query	Path	Edge Weights	Total Path Weight	Half	Explanation	Answer
`[1, 0]`	`1 → 0`	`[7]`	7	3.5	Sum from `1 → 0 = 7 >= 3.5`, median is node 0.	0
`[0, 1]`	`0 → 1`	`[7]`	7	3.5	Sum from `0 → 1 = 7 >= 3.5`, median is node 1.	1

Example 2:

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

Output: [1,0,2]

Explanation:

Query	Path	Edge Weights	Total Path Weight	Half	Explanation	Answer
`[0, 1]`	`0 → 1`	`[2]`	2	1	Sum from `0 → 1 = 2 >= 1`, median is node 1.	1
`[2, 0]`	`2 → 0`	`[4]`	4	2	Sum from `2 → 0 = 4 >= 2`, median is node 0.	0
`[1, 2]`	`1 → 0 → 2`	`[2, 4]`	6	3	Sum from `1 → 0 = 2 < 3`. Sum from `1 → 2 = 2 + 4 = 6 >= 3`, median is node 2.	2

Example 3:

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

Output: [2,2]

Explanation:

Query	Path	Edge Weights	Total Path Weight	Half	Explanation	Answer
`[3, 4]`	`3 → 1 → 0 → 2 → 4`	`[1, 2, 5, 3]`	11	5.5	Sum from `3 → 1 = 1 < 5.5`. Sum from `3 → 0 = 1 + 2 = 3 < 5.5`. Sum from `3 → 2 = 1 + 2 + 5 = 8 >= 5.5`, median is node 2.	2
`[1, 2]`	`1 → 0 → 2`	`[2, 5]`	7	3.5	Sum from `1 → 0 = 2 < 3.5`. Sum from `1 → 2 = 2 + 5 = 7 >= 3.5`, median is node 2.	2

Constraints:

2 <= n <= 10⁵
edges.length == n - 1
edges[i] == [u_i, v_i, w_i]
0 <= u_i, v_i < n
1 <= w_i <= 10⁹
1 <= queries.length <= 10⁵
queries[j] == [u_j, v_j]
0 <= u_j, v_j < n
The input is generated such that edges represents a valid tree.

Approaches

2 approaches with complexity analysis and trade-offs.

This approach handles each query independently by performing a full graph traversal to find the path between the two specified nodes. Once the path is found, it calculates the total weight and then iterates along the path to locate the weighted median node.

Algorithm

Build an adjacency list representation of the tree from the edges array.
For each query (u, v):
1. Perform a traversal (like BFS) starting from u to find the path to v. Use a map to store the parent of each visited node to allow path reconstruction.
2. Once v is found, trace back from v to u using the parent map to get the path.
3. Iterate through the path to calculate the sum of all edge weights, which is the total_weight.
4. Iterate again along the path from u to v, accumulating the edge weights in a current_weight variable.
5. The first node x (the destination of an edge) for which current_weight is at least half of total_weight is the median node for the query.

Walkthrough

For every query (u, v), we can find the unique path between them in the tree. A Breadth-First Search (BFS) or Depth-First Search (DFS) starting from u can be used to find v. During the traversal, we maintain a parent map to reconstruct the path once v is reached.

After finding v, we backtrack from v to u using the parent pointers to get the list of nodes on the path. Then, we calculate the total weight of this path by summing up the weights of its constituent edges. Finally, we traverse the path from u to v, keeping a running sum of the edge weights. The first node x where this running sum becomes greater than or equal to half of the total path weight is the answer for the query.

1import java.util.*;2 3class Solution {4    private List<int[]>[] adj;5 6    public int[] findWeightedMedian(int n, int[][] edges, int[][] queries) {7        adj = new ArrayList[n];8        for (int i = 0; i < n; i++) {9            adj[i] = new ArrayList<>();10        }11        for (int[] edge : edges) {12            adj[edge[0]].add(new int[]{edge[1], edge[2]});13            adj[edge[1]].add(new int[]{edge[0], edge[2]});14        }15 16        int[] ans = new int[queries.length];17        for (int i = 0; i < queries.length; i++) {18            ans[i] = solveQuery(queries[i][0], queries[i][1], n);19        }20        return ans;21    }22 23    private int solveQuery(int u, int v, int n) {24        if (u == v) return u;25 26        // 1. Find path from u to v using BFS and parent tracking27        Map<Integer, Integer> parentMap = new HashMap<>();28        Queue<Integer> queue = new LinkedList<>();29        queue.add(u);30        boolean[] visited = new boolean[n];31        visited[u] = true;32 33        while (!queue.isEmpty()) {34            int curr = queue.poll();35            if (curr == v) break;36            for (int[] neighbor : adj[curr]) {37                int nextNode = neighbor[0];38                if (!visited[nextNode]) {39                    visited[nextNode] = true;40                    parentMap.put(nextNode, curr);41                    queue.add(nextNode);42                }43            }44        }45 46        // 2. Reconstruct path47        List<Integer> path = new ArrayList<>();48        int curr = v;49        while (parentMap.containsKey(curr)) {50            path.add(curr);51            curr = parentMap.get(curr);52        }53        path.add(u);54        Collections.reverse(path);55 56        // 3. Calculate total path weight57        long totalWeight = 0;58        Map<Integer, Integer> weights = new HashMap<>();59        for (int i = 0; i < path.size() - 1; i++) {60            int node1 = path.get(i);61            int node2 = path.get(i + 1);62            int weight = getWeight(node1, node2);63            totalWeight += weight;64            weights.put(node2, weight);65        }66 67        // 4. Find median node68        long currentWeight = 0;69        for (int i = 1; i < path.size(); i++) {70            int node = path.get(i);71            currentWeight += weights.get(node);72            if (currentWeight * 2 >= totalWeight) {73                return node;74            }75        }76        77        return -1; // Should be unreachable for u != v78    }79 80    private int getWeight(int u, int v) {81        for (int[] neighbor : adj[u]) {82            if (neighbor[0] == v) {83                return neighbor[1];84            }85        }86        return -1; // Should not happen in a connected tree87    }88}

Complexity

Time

O(Q * N), where N is the number of nodes and Q is the number of queries. For each query, the traversal (BFS/DFS) to find the path can take up to O(N + E) = O(N) time in the worst case for a tree. Reconstructing and traversing the path also takes up to O(N) time.

Space

O(N) to store the adjacency list, the visited array, the parent map, and the path for each query.

Trade-offs

Pros

Simple to understand and implement.
Does not require complex data structures or algorithms.

Cons

Inefficient for a large number of queries, as it recomputes the path for each query from scratch.
Will result in a 'Time Limit Exceeded' error on platforms with large test cases.

This optimized approach preprocesses the tree to answer path-related queries efficiently. By pre-calculating parent pointers, depths, distances from the root, and a binary lifting table, we can find the Lowest Common Ancestor (LCA) and perform path calculations in logarithmic time per query.

Algorithm

Preprocessing:
1. Represent the tree using an adjacency list.
2. Perform a DFS from a root (e.g., node 0) to compute depth[i], parent[i], and dist[i] (weighted distance from root) for each node i.
3. Build a binary lifting table up[i][j] that stores the 2<sup>j</sup>-th ancestor of node i. This takes O(N log N) time.
Query Processing (for each query (u, v)):
1. Find the LCA, l = lca(u, v), in O(log N) time using the binary lifting table.
2. Calculate the total path weight: total_weight = dist[u] + dist[v] - 2 * dist[l].
3. Calculate half_weight = total_weight / 2.0.
4. If dist(u, l) >= half_weight, the median is on the u -> l path. Use binary lifting to find the node p just before the median on this path, and the answer is parent[p].
5. Otherwise, the median is on the l -> v path. Use binary lifting to find the median node directly by searching for the highest ancestor of v on the path from l that satisfies the median condition.

Walkthrough

The core idea is that the path between any two nodes u and v passes through their Lowest Common Ancestor (LCA). We can decompose the path into u -> ... -> lca(u,v) -> ... -> v.

Preprocessing:

Root the tree arbitrarily (e.g., at node 0).
Perform a single DFS traversal from the root to compute:
- depth[i]: The depth of each node i.
- parent[i]: The immediate parent of i in the rooted tree.
- dist[i]: The weighted distance from the root to i.
Use this information to build a binary lifting (or sparse) table, up[i][j], which stores the 2<sup>j</sup>-th ancestor of node i. This allows us to jump up the tree in powers of two.

Query Processing: For each query (u, v):

Find their LCA, l, in O(log N) time using the precomputed up table.
The total path weight is dist[u] + dist[v] - 2 * dist[l].
Calculate half_weight.
Check if the weighted distance from u to l (dist[u] - dist[l]) is greater than or equal to half_weight.
- If it is, the median node lies on the u -> l path. We can find it by binary lifting up from u.
- Otherwise, the median lies on the l -> v path. We can find it by binary lifting up from v. This binary lifting search for the median node also takes O(log N) time.

1import java.util.*;2 3class Solution {4    private List<long[]>[] adj;5    private int n;6    private int MAX_LOG;7    private int[] parent;8    private int[] depth;9    private long[] dist;10    private int[][] up;11 12    public int[] findWeightedMedian(int n, int[][] edges, int[][] queries) {13        this.n = n;14        this.MAX_LOG = (int) Math.ceil(Math.log(n) / Math.log(2));15        this.adj = new ArrayList[n];16        for (int i = 0; i < n; i++) adj[i] = new ArrayList<>();17        for (int[] edge : edges) {18            adj[edge[0]].add(new long[]{edge[1], edge[2]});19            adj[edge[1]].add(new long[]{edge[0], edge[2]});20        }21 22        this.parent = new int[n];23        this.depth = new int[n];24        this.dist = new long[n];25        this.up = new int[n][MAX_LOG];26 27        dfs(0, -1, 0, 0L);28 29        for (int j = 1; j < MAX_LOG; j++) {30            for (int i = 0; i < n; i++) {31                if (up[i][j - 1] != -1) {32                    up[i][j] = up[up[i][j - 1]][j - 1];33                } else {34                    up[i][j] = -1;35                }36            }37        }38 39        int[] ans = new int[queries.length];40        for (int i = 0; i < queries.length; i++) {41            int u = queries[i][0];42            int v = queries[i][1];43            if (u == v) {44                ans[i] = u;45                continue;46            }47 48            int l = lca(u, v);49            long totalWeight = dist[u] + dist[v] - 2 * dist[l];50            double halfWeight = totalWeight / 2.0;51            long dist_u_l = dist[u] - dist[l];52 53            if (dist_u_l >= halfWeight) {54                int curr = u;55                for (int j = MAX_LOG - 1; j >= 0; j--) {56                    if (up[curr][j] != -1 && depth[up[curr][j]] >= depth[l]) {57                        if (dist[u] - dist[up[curr][j]] < halfWeight) {58                            curr = up[curr][j];59                        }60                    }61                }62                ans[i] = parent[curr];63            } else {64                double targetDistFromLca = halfWeight - dist_u_l;65                int curr = v;66                for (int j = MAX_LOG - 1; j >= 0; j--) {67                    if (up[curr][j] != -1 && depth[up[curr][j]] >= depth[l]) {68                        if (dist[up[curr][j]] - dist[l] >= targetDistFromLca) {69                            curr = up[curr][j];70                        }71                    }72                }73                ans[i] = curr;74            }75        }76        return ans;77    }78 79    private void dfs(int u, int p, int d, long currentDist) {80        parent[u] = p;81        depth[u] = d;82        dist[u] = currentDist;83        up[u][0] = p;84        for (long[] edge : adj[u]) {85            int v = (int)edge[0];86            long w = edge[1];87            if (v != p) {88                dfs(v, u, d + 1, currentDist + w);89            }90        }91    }92 93    private int lca(int u, int v) {94        if (depth[u] < depth[v]) { int temp = u; u = v; v = temp; }95        96        for (int j = MAX_LOG - 1; j >= 0; j--) {97            if (up[u][j] != -1 && depth[u] - (1 << j) >= depth[v]) {98                u = up[u][j];99            }100        }101 102        if (u == v) return u;103 104        for (int j = MAX_LOG - 1; j >= 0; j--) {105            if (up[u][j] != -1 && up[v][j] != -1 && up[u][j] != up[v][j]) {106                u = up[u][j];107                v = up[v][j];108            }109        }110        return parent[u];111    }112}

Complexity

Time

O(N log N + Q log N). The preprocessing step, dominated by building the binary lifting table, takes O(N log N). Each query takes O(log N) for LCA calculation and another O(log N) for finding the median node.

Space

O(N log N) primarily for the `up` table used for binary lifting. The other arrays (`parent`, `depth`, `dist`) take O(N) space.

Trade-offs

Pros

Highly efficient and scalable for a large number of queries.
The preprocessing step allows for very fast query times.
A standard and powerful technique for solving path-related problems on trees.

Cons

More complex to implement compared to the brute-force approach.
Requires significant space for the binary lifting table.

Video walkthrough

Newsletter

One sharp idea, every week

System design and interview prep — short enough to finish.

No spam. Unsubscribe anytime.

Practice

Same difficulty — related problems to reinforce the pattern.

← Previous

Previous problem