Generated by Claude with claude-opus-4-7

Quick facts

Problem gist

Design a collection that supports two operations in any interleaving:

  • insert(num) — add a number to the collection.
  • find_kth_largest(k) — return the k-th largest element currently in the collection.

The twist versus the original LeetCode 703 is that k is a parameter of every query and may change from call to call. The classic min-heap-of-size-k solution relied on k being fixed at construction so it could safely discard everything below the running k-th largest. Once k can grow on a later call, no element is safe to throw away, and the whole strategy has to change.

The interesting part of this problem isn’t a single clever trick — it’s that there’s a family of valid data structures with very different insert/query tradeoffs, and the right choice depends on the workload.

Approaches

1. Sorted dynamic array — O(n) insert, O(1) query

Keep all values in one Python list, always in ascending order. On insert, binary-search for the position (O(log n)) and shift the suffix to open a slot (O(n) memory move). On query, arr[n - k] is the answer in O(1), because indexing into already-sorted contiguous memory is a single load.

The query cost is as low as it can physically get. The insert cost is the price.

2. Max-heap with on-demand selection — O(log n) insert, O(k log n) query

Push every value onto a max-heap as it arrives. The heap doesn’t know about k at all, so all n values are retained. To answer a query, pop the top k values; the last one popped is the k-th largest. Push them back so subsequent queries still see the whole collection.

This is the most natural extension of the original LC 703 idea. It pays at query time instead of insert time, and the payment scales with k, not with n. For small k (say, top-10 dashboards over a large stream), this beats anything that touches all n elements per query.

3. Balanced BST with subtree sizes — O(log n) insert, O(log n) query

The textbook answer. Use a self-balancing binary search tree (red-black, AVL, splay, or a randomized treap) and augment every node with the size of its subtree. Inserts follow the usual BST rules; subtree sizes are recomputed on the path from the new leaf back to the root, which is O(log n) because the tree is balanced.

Rank queries become a single root-to-leaf walk. Recast “k-th largest” as “(n - k + 1)-th smallest in in-order”. At each node, the left subtree contains all the in-order-smaller values, so:

  • if the target rank equals size(left) + 1, the current node is the answer;
  • if the target rank is ≤ size(left), descend left;
  • otherwise subtract size(left) + 1 from the rank and descend right.

The walk takes O(height) = O(log n) because the tree is balanced. This is the only approach that hits O(log n) for both operations, which is also the comparison-based lower bound — see the lower-bound follow-up below.

Among balanced BSTs, the randomized treap is by far the simplest to write from scratch: a BST keyed by value with each node carrying a random priority, plus the invariant that priorities form a max-heap top-down. Splits and merges propagate the priority order, so balance is maintained without any explicit rotation logic. Expected depth is O(log n). That’s what the code section below implements.

4. sortedcontainers.SortedList — what the library actually does

The Python idiom for this problem is from sortedcontainers import SortedList. It deserves explanation, because the surface API looks like a magic sorted array but the internal structure is what makes the asymptotics work.

SortedList keeps a list of sorted sublists (load factor 1000 by default, so each sublist holds at most ~1000 elements). Two pieces of metadata sit alongside it: a list of the first key of each sublist, and a cumulative-length index across sublists.

  • Insert: binary-search the sublist-key list to pick the right sublist (O(log(n / 1000))), then bisect.insort inside that sublist (an O(load) shift, but load is a small constant in practice). When a sublist exceeds the load factor it splits in half.
  • Index access: binary-search the cumulative-length index to find which sublist owns offset i (O(log(n / 1000))), then read the element at the local offset in O(1).

So both operations are O(log n) with very small constants, on cache-friendly contiguous arrays rather than scattered tree nodes. Asymptotically it matches the balanced BST; in practice it usually wins in CPython because pointer chasing through Python objects is slow. The data-structure family is essentially an unrolled O(sqrt(n)) skip list or a degenerate B-tree of height 2.

Choosing between them

ApproachInsertQueryBest when
Sorted dynamic arrayO(n)O(1)Queries vastly outnumber inserts (e.g., load once, query many)
Max-heap + extract kO(log n)O(k log n)Inserts dominate and k is small (top-10 over a large stream)
Balanced BST (treap)O(log n)O(log n)Mixed workload; want worst-case guarantees and a clean delete
SortedListO(log n)*O(log n)Same as BST, but you’re in Python — usually the practical winner

* amortized; sublist splits happen at a rate that keeps the load factor bounded.

The treap and SortedList share an asymptotic profile, but a hand-rolled treap is what’s expected in a whiteboard interview. SortedList is what you ship.

Python solution

Three independent implementations sharing the same API. Each is self-contained so they can be benchmarked or swapped without coordination.

import bisect
import heapq
import random
from typing import Optional


# =============================================================================
# Implementation 1: Sorted dynamic array
# =============================================================================
# Insert: O(n)   - bisect for position + shift the suffix.
# Query : O(1)   - direct indexed access.
# Right pick when queries vastly outnumber inserts.

class SortedArrayKthLargest:
    """Keeps all values in a contiguous list, always in ascending order."""

    def __init__(self) -> None:
        self._values: list[int] = []

    def insert(self, num: int) -> None:
        # bisect locates the insertion position in O(log n); the dominant cost
        # is the O(n) shift performed by list.insert under the hood.
        bisect.insort(self._values, num)

    def find_kth_largest(self, k: int) -> int:
        size = len(self._values)
        if not 1 <= k <= size:
            raise IndexError(f"k={k} out of range for collection of size {size}")
        # Ascending order: k-th largest sits at index size - k.
        return self._values[size - k]


# =============================================================================
# Implementation 2: Max-heap with on-demand selection
# =============================================================================
# Insert: O(log n)     - single heap push.
# Query : O(k log n)   - pop top k, return the k-th, restore the heap.
# Right pick when k is small relative to n.

class HeapKthLargest:
    """Stores all values in a max-heap, paying for the rank lookup at query time.

    Python's heapq is a min-heap, so values are stored negated; the smallest
    negative number at the root corresponds to the largest original value.
    """

    def __init__(self) -> None:
        self._max_heap_negated: list[int] = []

    def insert(self, num: int) -> None:
        heapq.heappush(self._max_heap_negated, -num)

    def find_kth_largest(self, k: int) -> int:
        size = len(self._max_heap_negated)
        if not 1 <= k <= size:
            raise IndexError(f"k={k} out of range for collection of size {size}")

        # Extract the top k values; the k-th pop is the k-th largest.
        extracted_top_k = [
            heapq.heappop(self._max_heap_negated) for _ in range(k)
        ]
        # Restore the heap so subsequent queries still see all n elements.
        for negated_value in extracted_top_k:
            heapq.heappush(self._max_heap_negated, negated_value)

        return -extracted_top_k[-1]


# =============================================================================
# Implementation 3: Order-statistics treap (balanced BST)
# =============================================================================
# Insert: expected O(log n)  - split + new node + merge.
# Query : expected O(log n)  - single root-to-leaf walk over subtree sizes.
#
# A treap is a BST keyed by value, heap-ordered by random priority. The random
# priority gives expected depth O(log n) without any explicit rebalancing logic
# (no rotations to maintain like in AVL/red-black). Each node also stores the
# size of its subtree so rank lookup is a single descent.


class _TreapNode:
    __slots__ = ("value", "priority", "left", "right", "subtree_size")

    def __init__(self, value: int) -> None:
        self.value = value
        self.priority = random.random()
        self.left: Optional["_TreapNode"] = None
        self.right: Optional["_TreapNode"] = None
        self.subtree_size = 1


def _subtree_size(node: Optional[_TreapNode]) -> int:
    return node.subtree_size if node is not None else 0


def _recompute_size(node: _TreapNode) -> None:
    node.subtree_size = 1 + _subtree_size(node.left) + _subtree_size(node.right)


def _split_by_value(
    node: Optional[_TreapNode], pivot: int
) -> tuple[Optional[_TreapNode], Optional[_TreapNode]]:
    """Split a treap into (values <= pivot, values > pivot).

    Returns two valid treaps. Both retain the BST and heap-priority invariants
    of their respective halves; subtree sizes are recomputed along the path.
    """
    if node is None:
        return None, None

    if node.value <= pivot:
        # Current node belongs on the left side; recurse into its right child
        # to peel off anything still greater than pivot.
        left_side, right_side = _split_by_value(node.right, pivot)
        node.right = left_side
        _recompute_size(node)
        return node, right_side

    # Current node belongs on the right side; recurse into its left child.
    left_side, right_side = _split_by_value(node.left, pivot)
    node.left = right_side
    _recompute_size(node)
    return left_side, node


def _merge_by_priority(
    left: Optional[_TreapNode], right: Optional[_TreapNode]
) -> Optional[_TreapNode]:
    """Merge two treaps where every value in `left` is <= every value in `right`.

    The merged tree is itself a valid treap: BST order is preserved by the
    precondition, and the higher-priority root is hoisted to keep the heap
    invariant top-down.
    """
    if left is None or right is None:
        return left if left is not None else right

    if left.priority > right.priority:
        left.right = _merge_by_priority(left.right, right)
        _recompute_size(left)
        return left

    right.left = _merge_by_priority(left, right.left)
    _recompute_size(right)
    return right


class TreapKthLargest:
    """Order-statistics treap with O(log n) insert and k-th largest queries."""

    def __init__(self) -> None:
        self._root: Optional[_TreapNode] = None

    def insert(self, num: int) -> None:
        # Split into values strictly less than num and values >= num, then
        # sandwich the new node between them. This admits duplicates cleanly.
        values_less_than_num, values_ge_num = _split_by_value(self._root, num - 1)
        new_leaf = _TreapNode(num)
        self._root = _merge_by_priority(
            _merge_by_priority(values_less_than_num, new_leaf),
            values_ge_num,
        )

    def find_kth_largest(self, k: int) -> int:
        size = _subtree_size(self._root)
        if not 1 <= k <= size:
            raise IndexError(f"k={k} out of range for collection of size {size}")

        # k-th largest in descending order = (size - k + 1)-th smallest in in-order.
        target_rank_from_smallest = size - k + 1
        return self._select_kth_smallest(self._root, target_rank_from_smallest)

    @staticmethod
    def _select_kth_smallest(node: Optional[_TreapNode], rank: int) -> int:
        """Iteratively descend, charging the rank against subtree sizes."""
        while node is not None:
            left_count = _subtree_size(node.left)
            if rank == left_count + 1:
                return node.value
            if rank <= left_count:
                node = node.left
            else:
                rank -= left_count + 1
                node = node.right
        raise RuntimeError("unreachable: rank validated before descent")

Interview follow-ups

What if k is fixed at construction time, like the original LeetCode 703?

When k is known up front and never changes, a min-heap of exactly k elements is optimal. Maintain the invariant that the heap holds the k largest values seen so far. On insert(num), push it; if the heap then has more than k elements, pop the smallest. The query is heap[0] in O(1).

Each insert is O(log k) and each query is O(1), using only O(k) space. This beats every dynamic-k structure listed above whenever k ≪ n, because it never stores values that cannot possibly be the answer. The reason this approach breaks when k becomes dynamic is exactly what makes it efficient here: it discards everything below the running k-th value, so a later query with a larger k has no way to retrieve them.

Can you achieve O(1) queries by trading off insert speed?

Yes — that’s exactly what the sorted dynamic array gives you. Keep all values in a contiguous sorted list, and a query becomes values[n - k]: a single indexed memory load, O(1). The trade is O(n) per insert because each one has to shift a tail of the array to open a slot.

The reason this doesn’t violate the Ω(log n) comparison lower bound is that the lower bound applies to searching for a rank in a comparison-based query — work you can’t avoid if the structure isn’t already laid out in rank order. The sorted array has already done that work, on every insert, in the form of memory shifts. Reads become free; writes pay the bill.

This is the right choice when queries vastly outnumber inserts. A monitoring system that loads a fixed dataset and then serves thousands of k-th percentile queries per second is the canonical case. The crossover with the BST approach happens around the point where insert frequency × n is comparable to query frequency × log n.

How would you support a delete(num) operation at the same asymptotic cost as insert?

Of the three implementations above, only the BST is naturally delete-friendly. In the treap, deletion splits the tree at num - 1 and at num, removes one occurrence from the middle piece (or all occurrences if you want set semantics), and merges the rest back together — three splits and two merges, all O(log n). SortedList exposes the same operation as remove(value) in O(log n).

The sorted array degrades to O(n) per delete (find then shift), matching its insert cost — fine if that’s already the chosen tradeoff. The heap is the awkward case: heapq has no delete-by-value primitive. The usual workaround is lazy deletion: track removed values in a side Counter, and when something surfaces at the top of the heap, check the counter and discard it if it’s been deleted. Amortized cost stays O(log n), but the implementation is significantly more error-prone and the heap can grow with deleted tombstones until enough surface to be pruned.

What if all values come from a small integer range U?

The comparison lower bound is a comparison-based bound. If values are drawn from a bounded integer range — say timestamps within a day, or page rankings 0–100 — you can break it.

Build a Fenwick tree (binary indexed tree) over value buckets, indexed [0, U). insert(num) increments the count at bucket num in O(log U). find_kth_largest(k) walks the Fenwick tree from the top bit downward, at each step checking whether the right half holds at least k items; if so, descend right, else descend left and subtract the right-half count from k. That’s an O(log U) rank-by-count query.

When U ≪ n, this is asymptotically faster than O(log n). It also handles duplicates naturally because Fenwick trees count, not store.

How does SortedList stay O(log n) amortized even though every insert shifts up to ~sqrt(n) elements within a sublist?

The trick is that the sublist size is bounded by a constant load factor (default 1000), not by sqrt(n). The “list-of-sublists” structure looks like sqrt decomposition, but SortedList keeps each sublist between roughly half the load factor and the full load factor by splitting full sublists in half. So the shift inside any one sublist is bounded by a constant in n, not by sqrt(n).

The O(log n) bound comes from elsewhere: the binary search over sublist boundaries and the index update. The shift inside the sublist is O(load) — which is a (large) constant. In the amortized argument, the split operation that happens occasionally is O(load) too, but it doubles the number of sublists only when one of them got fully populated, so the amortized cost per insert is O(log n) plus a small constant for the shift.

This is the same general technique B-trees use — keep node sizes bounded by a constant, pay an occasional split — but applied at one level instead of recursively.

How would you make this thread-safe for concurrent inserts and queries?

The simplest correct approach is a single threading.Lock acquired for the full duration of both insert and find_kth_largest. That serializes everything but is impossible to misuse, and for most workloads the contention will be acceptable.

For read-heavy workloads, a readers-writer lock (e.g., aiorwlock or a hand-rolled one over a threading.Condition) lets queries proceed in parallel while inserts hold an exclusive lock. This is a real win when queries outnumber inserts by orders of magnitude.

A different approach worth knowing: immutable snapshots with copy-on-write. Queries read a versioned snapshot of the BST; inserts copy the path from the new leaf to the root (O(log n) new nodes), leaving the rest shared. Old snapshots remain valid for in-flight queries. This is how persistent functional data structures work, and it gives lock-free reads at the cost of slightly more allocation per insert. The treap structure adapts to this cleanly because split/merge already builds new nodes for the touched path.

What’s the comparison-based lower bound, and is O(log n) tight for both operations?

For any data structure that supports both arbitrary inserts and rank queries using only comparisons between keys, Ω(log n) per operation is unavoidable. The argument is the same one that proves the Ω(n log n) sorting bound: with k comparisons you can distinguish at most 2^k outcomes, and there are n + 1 possible insertion positions or n possible rank answers, so k ≥ log₂(n).

The balanced BST hits this bound exactly for both operations, so it’s optimal in the comparison model. You can do better only by leaving that model — using value-range tricks (Fenwick tree over buckets, van Emde Boas trees with O(log log U) per operation), assuming a specific distribution, or accepting approximation (t-digest, KLL sketches, which give O(log n) rank queries with bounded error in sublinear space).