# Find Weighted Median Node in Tree **Difficulty:** HARD [External](https://leetcode.com/problems/find-weighted-median-node-in-tree) Canonical: https://scaleengineer.com/dsa/problems/find-weighted-median-node-in-tree **Patterns:** [Dynamic Programming](https://scaleengineer.com/dsa/patterns/dynamic-programming) **Algorithms:** [Binary Search](https://scaleengineer.com/algorithms/binary-search), [Depth-First Search](https://scaleengineer.com/algorithms/depth-first-search) **Data structures:** Array, Tree --- ## Problem 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] = [ui, vi, wi]` indicates an edge from node `ui` to `vi` with weight `wi`. The **weighted median node** is defined as the **first** node `x` on the path from `ui` to `vi` such that the sum of edge weights from `ui` 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] = [uj, vj]`, determine the weighted median node along the path from `uj` to `vj`. 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:** ![](https://assets.glich.co/dsa/find-weighted-median-node-in-tree/image0.png) | Query | Path | EdgeWeights | TotalPathWeight | 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\] **E** **xplanation:** ![](https://assets.glich.co/dsa/find-weighted-median-node-in-tree/image1.png) | Query | Path | EdgeWeights | TotalPathWeight | 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:** ![](https://assets.glich.co/dsa/find-weighted-median-node-in-tree/image2.png) | Query | Path | EdgeWeights | TotalPathWeight | 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 <= 105` * `edges.length == n - 1` * `edges[i] == [ui, vi, wi]` * `0 <= ui, vi < n` * `1 <= wi <= 109` * `1 <= queries.length <= 105` * `queries[j] == [uj, vj]` * `0 <= uj, vj < n` * The input is generated such that `edges` represents a valid tree. # Approaches ## Brute-Force Path Traversal for Each Query 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. **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. **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. ### Explanation 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. ```java import java.util.*; class Solution { private List[] adj; public int[] findWeightedMedian(int n, int[][] edges, int[][] queries) { adj = new ArrayList[n]; for (int i = 0; i < n; i++) { adj[i] = new ArrayList<>(); } for (int[] edge : edges) { adj[edge[0]].add(new int[]{edge[1], edge[2]}); adj[edge[1]].add(new int[]{edge[0], edge[2]}); } int[] ans = new int[queries.length]; for (int i = 0; i < queries.length; i++) { ans[i] = solveQuery(queries[i][0], queries[i][1], n); } return ans; } private int solveQuery(int u, int v, int n) { if (u == v) return u; // 1. Find path from u to v using BFS and parent tracking Map parentMap = new HashMap<>(); Queue queue = new LinkedList<>(); queue.add(u); boolean[] visited = new boolean[n]; visited[u] = true; while (!queue.isEmpty()) { int curr = queue.poll(); if (curr == v) break; for (int[] neighbor : adj[curr]) { int nextNode = neighbor[0]; if (!visited[nextNode]) { visited[nextNode] = true; parentMap.put(nextNode, curr); queue.add(nextNode); } } } // 2. Reconstruct path List path = new ArrayList<>(); int curr = v; while (parentMap.containsKey(curr)) { path.add(curr); curr = parentMap.get(curr); } path.add(u); Collections.reverse(path); // 3. Calculate total path weight long totalWeight = 0; Map weights = new HashMap<>(); for (int i = 0; i < path.size() - 1; i++) { int node1 = path.get(i); int node2 = path.get(i + 1); int weight = getWeight(node1, node2); totalWeight += weight; weights.put(node2, weight); } // 4. Find median node long currentWeight = 0; for (int i = 1; i < path.size(); i++) { int node = path.get(i); currentWeight += weights.get(node); if (currentWeight * 2 >= totalWeight) { return node; } } return -1; // Should be unreachable for u != v } private int getWeight(int u, int v) { for (int[] neighbor : adj[u]) { if (neighbor[0] == v) { return neighbor[1]; } } return -1; // Should not happen in a connected tree } } ``` ### 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. ## Lowest Common Ancestor (LCA) with Binary Lifting 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. **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. **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. ### Explanation 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:** 1. Root the tree arbitrarily (e.g., at node 0). 2. 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`. 3. Use this information to build a binary lifting (or sparse) table, `up[i][j]`, which stores the 2j-th ancestor of node `i`. This allows us to jump up the tree in powers of two. **Query Processing:** For each query `(u, v)`: 1. Find their LCA, `l`, in O(log N) time using the precomputed `up` table. 2. The total path weight is `dist[u] + dist[v] - 2 * dist[l]`. 3. Calculate `half_weight`. 4. 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. ```java import java.util.*; class Solution { private List[] adj; private int n; private int MAX_LOG; private int[] parent; private int[] depth; private long[] dist; private int[][] up; public int[] findWeightedMedian(int n, int[][] edges, int[][] queries) { this.n = n; this.MAX_LOG = (int) Math.ceil(Math.log(n) / Math.log(2)); this.adj = new ArrayList[n]; for (int i = 0; i < n; i++) adj[i] = new ArrayList<>(); for (int[] edge : edges) { adj[edge[0]].add(new long[]{edge[1], edge[2]}); adj[edge[1]].add(new long[]{edge[0], edge[2]}); } this.parent = new int[n]; this.depth = new int[n]; this.dist = new long[n]; this.up = new int[n][MAX_LOG]; dfs(0, -1, 0, 0L); for (int j = 1; j < MAX_LOG; j++) { for (int i = 0; i < n; i++) { if (up[i][j - 1] != -1) { up[i][j] = up[up[i][j - 1]][j - 1]; } else { up[i][j] = -1; } } } int[] ans = new int[queries.length]; for (int i = 0; i < queries.length; i++) { int u = queries[i][0]; int v = queries[i][1]; if (u == v) { ans[i] = u; continue; } int l = lca(u, v); long totalWeight = dist[u] + dist[v] - 2 * dist[l]; double halfWeight = totalWeight / 2.0; long dist_u_l = dist[u] - dist[l]; if (dist_u_l >= halfWeight) { int curr = u; for (int j = MAX_LOG - 1; j >= 0; j--) { if (up[curr][j] != -1 && depth[up[curr][j]] >= depth[l]) { if (dist[u] - dist[up[curr][j]] < halfWeight) { curr = up[curr][j]; } } } ans[i] = parent[curr]; } else { double targetDistFromLca = halfWeight - dist_u_l; int curr = v; for (int j = MAX_LOG - 1; j >= 0; j--) { if (up[curr][j] != -1 && depth[up[curr][j]] >= depth[l]) { if (dist[up[curr][j]] - dist[l] >= targetDistFromLca) { curr = up[curr][j]; } } } ans[i] = curr; } } return ans; } private void dfs(int u, int p, int d, long currentDist) { parent[u] = p; depth[u] = d; dist[u] = currentDist; up[u][0] = p; for (long[] edge : adj[u]) { int v = (int)edge[0]; long w = edge[1]; if (v != p) { dfs(v, u, d + 1, currentDist + w); } } } private int lca(int u, int v) { if (depth[u] < depth[v]) { int temp = u; u = v; v = temp; } for (int j = MAX_LOG - 1; j >= 0; j--) { if (up[u][j] != -1 && depth[u] - (1 << j) >= depth[v]) { u = up[u][j]; } } if (u == v) return u; for (int j = MAX_LOG - 1; j >= 0; j--) { if (up[u][j] != -1 && up[v][j] != -1 && up[u][j] != up[v][j]) { u = up[u][j]; v = up[v][j]; } } return parent[u]; } } ``` ### 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 2j-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.