Kth Largest Element, Two Ways

The Kth largest question asks for a single value: the element at position k if the list were sorted in descending order. Duplicates occupy separate positions, so in [3, 2, 3, 1, 2, 4, 5, 5, 6] the 4th largest is 4, reached through 6, 5, 5, 4. Ask about duplicates in the first minute, because the 4th largest distinct value is a different question, and its answer is 3.

Way 1: the capped heap's root is the answer

The top-K heap from the previous section already holds this answer. Once every value has passed through a min-heap capped at k, the root is the smallest of the k largest values, which is the Kth largest by definition. Seeding the heap with heapify on the first k values replaces k pushes with a single O(k) build.

import heapq

def kth_largest_heap(nums, k):
    if not 1 <= k <= len(nums):
        raise ValueError("k out of range")
    heap = nums[:k]
    heapq.heapify(heap)
    for x in nums[k:]:
        if x > heap[0]:
            heapq.heapreplace(heap, x)
    return heap[0]

print(kth_largest_heap([3, 2, 3, 1, 2, 4, 5, 5, 6], 4))

Without the range check, k = 20 on these 9 values builds a heap of all 9 and returns their minimum, 1, with no error. A wrong number that looks plausible is harder to catch than a traceback, so raise.

Way 2: quickselect partitions in place

Quickselect borrows the partition step from quicksort. Pick a pivot, move every smaller value to its left, and drop the pivot into the slot it would occupy in ascending order. If that slot is index n - k, the pivot is the answer. Otherwise the answer lies on a single side, so the loop discards the other side and partitions again, touching a shrinking range each pass.

def kth_largest_select(nums, k):
    target = len(nums) - k
    lo, hi = 0, len(nums) - 1
    while True:
        pivot = nums[hi]
        store = lo
        for i in range(lo, hi):
            if nums[i] < pivot:
                nums[store], nums[i] = nums[i], nums[store]
                store += 1
        nums[store], nums[hi] = nums[hi], nums[store]
        print(pivot, store, nums)
        if store == target:
            return nums[store]
        if store < target:
            lo = store + 1
        else:
            hi = store - 1

data = [3, 2, 1, 5, 6, 4]
print(kth_largest_select(data, 2))

4 3 [3, 2, 1, 4, 6, 5]
5 4 [3, 2, 1, 4, 5, 6]
5

The target index is 6 - 2 = 4. The first pass uses pivot 4, which lands at index 3, so the answer lies to the right and lo moves to 4. The second pass partitions only indexes 4 and 5: pivot 5 lands at index 4, which is the target, and the function returns 5. The list data came back reordered as [3, 2, 1, 4, 5, 6], because quickselect mutates its input; pass nums[:] when the caller still needs the original order.

Choosing the last element as pivot keeps the animation repeatable and creates a weakness. On input already sorted ascending, the last element is always the largest in range, so each pass removes a single element whenever the target lies to its left, and the work grows to O(n^2). Swapping a random element into position hi before each pass, chosen with random.randint(lo, hi), makes that pattern vanishingly unlikely and bounds the expected cost at about 3.4n comparisons for any k.

Choosing between them

The heap's O(k) is all the memory it ever needs. Quickselect's O(1) assumes all n values are already loaded, which a stream never provides.

Put numbers on it. With n = 1,000,000 and k = 100, the heap costs about n log2 k, roughly 6.6 million comparisons in the worst case, and holds 100 values. Quickselect expects at most about 3.4 million comparisons but needs all 1,000,000 values in memory and permission to reorder them. Values already sitting in a Python list on one machine favor quickselect. Values arriving from a stream, or a k tiny next to n, favor the heap. For a single query on a short list, sorted(nums, reverse=True)[k - 1] is the readable choice and a fine answer.

About this lesson

This section is part of the Heap & Top-K: Beginner lesson on DataDriven, the home of data engineering. Lessons cover SQL, Python, Spark, data modeling and pipeline architecture, and each section pairs its explanation with worked examples and practice you run as you go.