p18 — Kth Smallest Element in a BST

Source: LeetCode 230 · Medium · Topics: Tree, DFS, BST, Inorder Companies (2024–2025 frequency): Amazon (very high), Meta (high), Google (high), Microsoft (high), Bloomberg (medium), TikTok (medium) Loop position: mid-loop algorithmic round; ALWAYS comes with the augmented-BST follow-up (“what if there are frequent insertions/deletions?”). The follow-up is the real test.

1. Quick Context

The base problem (return the k-th smallest value in a BST) is a 10-line iterative inorder with early termination — straightforward once you internalize the inorder template from p16. The follow-up is what every interviewer is actually waiting for: “Suppose the BST is modified often (frequent inserts and deletes), and you need to find the k-th smallest many times. How would you optimize?”

The intended answer: augment each node with the size of its subtree (node.size = 1 + size(left) + size(right)). Then kth_smallest(node, k):

• left_size = size(node.left)
• If k <= left_size → recurse left with same k.
• If k == left_size + 1 → this node IS the answer.
• Else → recurse right with k - left_size - 1.

Each query becomes O(H). Insert/delete maintain size in O(H). This order-statistic tree pattern is the foundation of every “find rank / select k-th” algorithm in production databases.

What it looks like it tests: can you walk inorder. What it actually tests: (a) do you stop inorder early at the k-th element (memory-efficient streaming) and (b) do you know order-statistic trees / augmented BSTs.

2. LeetCode Link + Attempt Gate

🔗 https://leetcode.com/problems/kth-smallest-element-in-a-bst/

STOP. 7-min timer. Before peeking at hints, also pre-think the augmentation follow-up — it’s coming.

3. Prerequisite Concepts

• p16 — iterative inorder template (you should be able to write it cold)
• p17 — descending the BST exploiting order
• The notion of “k-th order statistic”
• Subtree size aggregation (basic tree DP)

4. How to Approach (FRAMEWORK Steps 1–9)

Step 1 — Restate: “Given the root of a BST and an integer k (1-indexed), return the k-th smallest value in the BST.”

Step 2 — Clarify:

• “1-indexed or 0-indexed?” (LC: 1-indexed.)
• “Is k guaranteed valid (1 ≤ k ≤ N)?” (LC: yes.)
• “Duplicates?” (LC: BST guarantees uniqueness here.)
• “Will the tree be modified between queries?” (Crucial clarifying question — pre-empts the follow-up. If the answer is “yes,” skip to augmented BST.)
• “Single query or many?” (Same direction as previous.)

Step 3 — Constraints: N up to 10^4. k between 1 and N. Need O(H + k) for single query, O(log N + log N) per query for augmented version.

Step 4 — Examples:

• [3,1,4,null,2], k=1 → 1.
• [5,3,6,2,4,null,null,1], k=3 → 3.
• Single node, k=1 → that node’s value.
• Right-skewed tree of 10 nodes, k=5 → the 5th value.

Step 5 — Brute Force: Collect all values via full inorder into a list, return list[k-1]. O(N) time, O(N) space. State this then propose early-termination.

Step 6 — Complexity (optimal):

• Single query, unaugmented: O(H + k) time, O(H) space.
• Many queries, augmented: O(log N) per query (balanced), O(H) per update.

Step 7 — Pattern Recognition:

• “k-th smallest in an ordered tree” → early-terminating inorder.
• “many queries with mutations” → augment with subtree size (order-statistic tree). This is the same pattern as Fenwick-tree-of-counts and segment-tree-of-sizes from Phase 3.

Step 8 — Develop:

• Iterative inorder with counter: descend leftmost, pop, decrement k, if k == 0 return val, else descend right. Stops as soon as k-th is popped.
• Recursive inorder with closure counter: same, with nonlocal k and early return via flag.
• Augmented BST: maintain size on every node; descend based on size(left) vs k.

Step 9 — Prove correctness (iterative inorder + counter): Inorder visits values in strictly increasing order (BST property). The counter decrements by 1 on each visit and the k-th visit produces the value with rank exactly k (1-indexed). The walk stops immediately, so time is O(H) descent + O(k) pops = O(H + k). ∎

5. Progressive Hints

If stuck for 7 min, open hints.md. The early-termination is the easy part; the augmented-BST follow-up is where you should spend hint #5.

6. Deeper Insight — Order-Statistic Trees

The augmented BST technique is one of the most reusable data-structure ideas in interviews:

Add ONE extra field to each node whose value depends only on the subtree, AND can be maintained in O(1) per node during inserts/deletes (by propagating up the modified path). This converts queries that would require O(N) scans into O(H) walks.

Examples of this pattern:

• Subtree size → k-th smallest, rank of an element, “count of values < x.”
• Subtree sum → range sum on indexed data.
• Subtree max/min → range max/min queries.
• Subtree height → AVL tree balancing.
• Black-node count → Red-Black tree invariants.

The cost is one extra field and ~3 extra lines in insert/delete. The win is converting linear-scan queries into logarithmic ones. This is the augmentation argument — internalize it; it appears in Phase 3 (segment trees, Fenwick) and in every database index design.

7. Anti-Pattern Analysis

1. Full inorder into a list, then list[k-1]. Correct but wastes O(N) time and space. State as brute force; switch to early-terminating version.
2. Recursive inorder that doesn’t terminate early. Decrementing k works, but return from a recursive helper that already finished doesn’t stop the rest of the recursion unless you check a flag at each entry. Beginners write code that finds the answer but continues to walk the rest of the tree — fine for correctness, embarrassing for performance.
3. Off-by-one with 0-indexed vs 1-indexed. LC uses 1-indexed. Always restate.
4. Augmenting size as “number of nodes in left subtree” instead of “total subtree size including this node.” Both are valid conventions but you must be consistent. “Total” is more common in textbooks; “left only” is more common in some implementations. State your choice in a comment.
5. For the augmented follow-up: forgetting that you must maintain size on EVERY insert/delete. A common bug: implementing the query correctly but not updating sizes on mutation. The next query returns garbage. Always pair augmentation logic with the mutation code that maintains it.

8. Skills & Takeaways

• Streaming traversal with early termination — the iterative inorder template is reusable across LC 230, 285, 99, 173.
• Order-statistic trees as a named pattern. Cite it by name in the follow-up — it signals you know the textbook (“Introduction to Algorithms” Ch. 14).
• Augmentation discipline: any extra field must be maintained in EVERY mutation. State this trade-off explicitly.
• Query/mutation balance: know when to invest preprocessing cost (augmentation) vs query cost (linear scan).

9. When to Move On

Move on when:

• You write iterative inorder with k-counter early termination from memory in <3 minutes.
• You can sketch the augmented-BST follow-up (size field + descent logic) on a whiteboard in <5 minutes.
• You can articulate the trade-off: “unaugmented = O(H+k) per query, no overhead on inserts; augmented = O(log N) per query, +O(1) bookkeeping per insert/delete.”
• Your stress test confirms agreement with a sorted-list oracle on 1000 random BSTs.

10. Company Context

Meta:

• This is a Meta favorite. The base problem is the warm-up; the augmented follow-up is the real interview. Often 30 min on follow-up, 10 on base.
• Watch for “implement the insert/delete that maintains size” — make sure you can do it.
• Scorecard phrase: “Algorithm — recognized augmentation pattern; implemented insert with O(1) size maintenance.”

Google:

• Will push on proof of correctness for the augmented descent (“prove that recursing left with the same k is correct when k ≤ size(left)”) and on average vs worst case complexity.
• Scorecard phrase: “Deep Dive — order-statistic tree; clean proof; honest about balancing requirement for log guarantees.”

Amazon:

• Common SDE-II round. Look for: “what if this is a billion-element tree on disk?” → augment a B-tree; same descent logic with size per B-tree node.
• Scorecard phrase: “Coding Excellence — clean iterative inorder; extended to large-scale via B-tree augmentation.”

Microsoft:

• May ask the kth largest version — easy modification: reverse-inorder (right-root-left).
• Scorecard phrase: “Problem Solving — generalized to kth largest; recognized symmetry.”

Bloomberg:

• Frames as “find the kth-most-traded stock in our order book” — order book BSTs are augmented with volume/size in production.
• Scorecard phrase: “Strong — connected to real Bloomberg data structures.”

TikTok / ByteDance:

• Often pairs with LC 295 (Find Median from Data Stream) in the same loop — both are “find rank/order-statistic in dynamic data” problems.
• Scorecard phrase: “Solid — recognized the order-statistic family across problems.”

11. Interviewer’s Lens

12. Level Delta

Mid (L4):

• Writes recursive inorder solution after prompting for early termination.
• Discusses augmentation when explicitly asked.

Senior (L5):

• Writes iterative inorder with k-counter first try.
• Proactively raises the “many queries with updates” follow-up before being asked.
• Implements size augmentation correctly including the insert path.

Staff (L6):

• Frames the problem upfront as “this is asking about order statistics — am I optimizing for query or update?”
• Implements augmented BST cleanly; discusses why AVL/RB are needed for log-time worst case.
• Connects to Fenwick tree (for static order statistics on indexed data) and to database B+tree indexing.

Principal (L7+):

• Discusses concurrent order-statistic queries on a shared, mutating tree (snapshot isolation, MVCC).
• Identifies the meta-pattern: “augment with a subtree summary that is constant-time mergeable” — same idea drives segment-tree lazy propagation.
• Sketches the design of a database column index that supports rank/order queries in production (e.g., PostgreSQL B-tree with covering indexes).

13. Follow-up Q&A

Q1: What if the BST is modified frequently and you need many k-th smallest queries? A: Augment each node with size = 1 + size(left) + size(right). Maintain on every insert/delete (size updates propagate up the modified path in O(H)). Then kth_smallest(node, k): let L = size(node.left). If k <= L, recurse left. If k == L+1, return node.val. Else recurse right with k - L - 1. Each query is O(H).

Q2: What if the tree could be unbalanced — what’s the worst-case? A: H = N for fully skewed BST (sorted insertion order). To get O(log N) worst-case, use a self-balancing variant (AVL, RB-tree, or Treap). The augmentation pattern is identical; only the rebalancing logic changes.

Q3: How would you find the rank of a given value (the inverse query)? A: With augmentation: descend from root. If value < node.val, recurse left. If value > node.val, accumulate size(node.left) + 1 and recurse right. If value == node.val, return accumulated + size(node.left) + 1. O(H).

Q4: Find the median (the n/2-th smallest)? A: Same as kth-smallest with k = ceil(N/2). Requires knowing N — keep a tree-level counter or use root’s size if augmented.

Q5: How would you do this on a B-tree (each node has up to M children)? A: Augment each child pointer with the subtree-size below it. To find k-th: at each node, scan children left-to-right accumulating sizes; recurse into the first child whose cumulative size ≥ k, adjusting k. O(log_M N · M) per query.

14. Solution Walkthrough

See solution.py:

• kth_smallest_iterative(root, k) — iterative inorder with k-counter early termination (the shippable form).
• kth_smallest_recursive(root, k) — recursive inorder with closure-captured counter + early flag.
• kth_smallest_brute_full_inorder(root, k) — full inorder into list (oracle for stress tests).
• OrderStatBST class — augmented BST with size field. Methods: insert, delete, kth_smallest, rank (the follow-up implementation).
• stress_test() — Part A: 1000 random BSTs; for each, run all k from 1..N, assert all 3 unaugmented variants agree. Part B: 1000 random insert/delete/query sequences against the augmented OrderStatBST, cross-check with a sorted-list oracle.
• edge_cases() — single node, deepest k, left-skewed, right-skewed, augmented insert-delete-insert correctness.

Run: python3 solution.py → “ALL TESTS PASSED” on stdout.

15. Beyond the Problem

Principal-engineer code review comment:

“Base kth-smallest looks fine, ship it. About the augmented version — two flags. First, your delete rebalances size on the deletion path but you’re using simple BST delete without rotations, so you’ve added O(1) bookkeeping per node touched but the tree itself degrades on adversarial input. If the spec promises log-time worst case you need AVL or RB rotations, and those rotations themselves need to re-augment size on the rotated nodes — that’s a 3-line change but if you forget it, queries silently return wrong answers for hours before someone notices. Add a unit test that calls kth_smallest(k=mid) AFTER n/2 random insertions and n/4 random deletions, asserting against a sorted oracle. Second, the size field is an int; for a billion-node tree you’re fine, but for a serialized format please pick uint64 explicitly because Python’s int is unbounded and that hides bugs that bite when ported to typed languages.”

Connections forward:

• p19 (Sorted Array → BST) builds a balanced BST — the natural pair with this problem.
• Phase 3 Fenwick trees and segment trees are augmentation patterns on arrays (vs trees here).
• LC 173 (BST Iterator) is the streaming version of the iterative inorder template here — same code, different interface.
• LC 315 (Count of Smaller Numbers After Self) uses an order-statistic BST as one valid solution.
• Database B+tree indexes in PostgreSQL/MySQL use this augmentation in production.