Skip to main content
Talentd

Find Weighted Median Node in Tree - Solution & Explanation

HardArrayBinary SearchDynamic ProgrammingBit Manipulation21 min readAsked at: Google
Practice this problem

Problem Statement

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:

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

Approach Overview

Problem Overview: You are given a tree where each node has an associated weight. The goal is to find a weighted median node: a node such that if you remove it, every resulting connected component has total weight ≤ half of the total weight of the entire tree.

Approach 1: Brute Force Component Check (O(n^2) time, O(n) space)

Try every node as the candidate median. For each node, temporarily treat it as removed and run a DFS into each adjacent subtree to compute the total weight of that component. Track the largest component weight. If the maximum component weight is ≤ totalWeight / 2, that node satisfies the weighted median condition. This approach is simple and directly models the definition, but each candidate requires traversing large portions of the tree again. With n nodes, this leads to O(n^2) time in the worst case.

Approach 2: Subtree Weight DFS (Optimal) (O(n) time, O(n) space)

Compute subtree weights using a single depth-first search. Let subtree[u] represent the total weight of node u and all nodes below it. After computing these values, evaluate each node as a candidate median. For node u, the weight of every child component is simply subtree[child]. The remaining component (the "parent side") has weight totalWeight - subtree[u]. The largest of these component weights determines whether u is valid. If all components are ≤ totalWeight / 2, you found the weighted median node.

This works because subtree sums allow constant-time evaluation of each component after a single traversal. The tree is processed once to compute weights and once more to check the condition, giving linear time overall. The approach relies heavily on tree structure properties and precomputed subtree aggregates, a common pattern in dynamic programming on trees.

Recommended for interviews: The subtree-weight DFS approach. Interviewers expect you to recognize that recomputing component weights repeatedly is wasteful. Precomputing subtree sums reduces the complexity from O(n^2) to O(n). Starting with the brute-force idea shows you understand the definition, but optimizing it with DFS aggregation demonstrates strong tree algorithm skills.

Solution

Thinking

The weighted median is the first vertex on u \to v whose prefix weight from u is at least half the path. n,q \le 10^5 call for LCA and weighted prefixes.

After lca and the total W, binary-lift on u \to lca or lca \to v to the farthest node whose prefix is still < W/2, then take one more step.

Root the tree at node 0. Use BFS to compute each node's depth depth, parent p, and weighted distance dist from the root, and build the binary lifting table f[i][j] as the 2^j-th ancestor of i.

For a query (u, v): if u = v, the answer is u. Otherwise let x = lca(u, v) and W = dist[u] + dist[v] - 2 cdot dist[x]. The weighted median is the first node on the path starting from u whose prefix weight is at least W / 2. We compare 2 cdot pref \ge W to avoid floating-point arithmetic.

  • If 2 cdot (dist[u] - dist[x]) \ge W, the median lies on u \to x (including x). Lift from u to the farthest ancestor k that still satisfies 2 cdot (dist[u] - dist[k]) < W, then take one parent step to p[k].
  • Otherwise the median lies on x \to v (excluding x). Lift from v to the highest node whose depth is greater than x and 2 cdot (dist[u] + dist[k] - 2 cdot dist[x]) \ge W.

The time complexity is O((n + q) times log n), and the space complexity is O(n times log n), where n is the number of nodes and q is the number of queries.

Code

PythonJavaC++GoTypeScriptRustC#

Python

Java

C++

Go

TypeScript

Rust

C#

Try this approach in the editor →

Detailed Complexity Analysis

ApproachTimeSpaceWhen to Use
Brute Force Component DFSO(n^2)O(n)Useful for understanding the definition of weighted median or when constraints are very small
Subtree Weight DFS (Optimal)O(n)O(n)General solution for large trees; precompute subtree sums and check each node efficiently
Tree DP / Rerooting VariantO(n)O(n)Useful when the problem extends to multiple queries or requires evaluating properties from different roots

Video Solution

3585. Find Weighted Median Node in Tree | LeetCode weekly contest 454 • Amit Dhyani • 1,207 views views

Watch 4 more video solutions →

Frequently Asked Questions

Is Find Weighted Median Node in Tree easy or hard?
Find Weighted Median Node in Tree is typically classified as a hard problem. The difficulty comes from recognizing the centroid-like property and avoiding repeated subtree computations. Efficient solutions require strong understanding of DFS and tree dynamic programming techniques.
Find Weighted Median Node in Tree Python/Java solution
Implement a DFS to compute subtree weights and store them in an array. Then iterate through nodes and evaluate component weights using the precomputed values. The same logic works in Python, Java, C++, and Go with an adjacency list and recursive DFS.
How to solve Find Weighted Median Node in Tree in O(n)?
First compute the total subtree weight for every node using a depth-first search. Then for each node, check the weight of every child subtree and the remaining part of the tree (totalWeight - subtree[node]). If the largest component weight is at most half the total weight, that node is the weighted median. Because subtree sums are precomputed, each check is constant time.
What is the best approach for Find Weighted Median Node in Tree?
The optimal approach uses a DFS to compute subtree weight sums. After calculating subtree[u] for every node, evaluate each node by checking the weight of its child subtrees and the remaining parent component. If every component has weight ≤ totalWeight/2, that node is the weighted median. This runs in O(n) time and O(n) space.
Is Find Weighted Median Node in Tree asked at Google/Amazon/Meta?
Weighted median and centroid-style tree problems appear frequently in interviews at companies like Google, Amazon, and Meta. They test understanding of tree traversal, subtree aggregation, and component analysis after removing a node. Variants often appear in graph and tree algorithm interview rounds.
What data structure is used in Find Weighted Median Node in Tree?
The core data structure is an adjacency list representing the tree. A depth-first search traverses the structure to compute subtree sums. Arrays or vectors store subtree weights, while recursion or an explicit stack performs the traversal.
What is the time complexity of Find Weighted Median Node in Tree?
The optimal DFS-based solution runs in O(n) time where n is the number of nodes in the tree. One traversal computes subtree weights, and another pass checks the weighted median condition for each node. Space complexity is O(n) due to the adjacency list and recursion stack.

Ready to solve this problem?

Practice Find Weighted Median Node in Tree with our built-in code editor and test cases.

Practice on FleetCode
Previous

Two Sum

Next

Median of Two Sorted Arrays