Finding the Kth Largest Element

There are three main options and they differ in what they spend. Sorting the array and indexing into it costs O of n log n time and no extra space. Keeping a heap of the k largest elements seen so far costs O of n log k time and O of k space. Quickselect partitions the array like quicksort but recurses into only one side, giving O of n average time with O of 1 space.

def kth_largest(values, k):
    ordered = sorted(values, reverse=True)
    return ordered[k - 1]


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

The outputs are 5 and 1. The second largest of that array is 5, and the smallest, which is the sixth largest, is 1.

The mechanism of the sorting version is trivial and it is worth stating why it is not automatically the wrong answer. It is linear extra work, not quadratic, and for a single query on a moderate array it is almost always fast enough. The real issue is that the bound scales badly with n, and it reorders the input as a side effect unless you sort a copy.

The trade-off to name is what the interview is really testing. If k is small and the array is huge, the heap version wins because it never stores more than k items and its work per element is logarithmic in k rather than in n. If you will run many queries against the same array, sorting once and keeping the result is best. If you need the worst case to be good, quickselect with a random pivot or the median-of-three heuristic is the one to pick, because plain quickselect degrades to quadratic on adversarial input, and the heap version has a guaranteed bound if you prefer predictability.

Python's heapq module gives you the two operations you need: heappush to add an item and heapreplace to swap the root with a new item and restore the heap. Keeping only the k largest elements is enough, because anything smaller than the kth best can never make the answer.

import heapq

data = [3, 2, 1, 5, 6, 4]
largest = heapq.nlargest(3, data)
print(largest)
print(heapq.nsmallest(2, data))

The outputs are 6, 5, 4 and 1, 2. The first call returns the three largest values in descending order and the second returns the two smallest in ascending order.

The mechanism is the heap invariant. A min heap keeps the smallest element at the root, so if the heap holds the k largest items seen so far, the root is the kth largest so far. When a new value arrives, if it is bigger than the root it belongs in the heap and the root is discarded; otherwise it is ignored. Either way the heap never grows beyond k items, and each step costs O of log k.

The trade-off to name is that the bound is O of n log k rather than O of n. When k is small that is close to linear, and the memory use is bounded by k regardless of how large the input is, which is exactly what you want for a streaming scenario where the input never fits in memory. The subtle part is that this is only true if you implement it with a fixed size heap. The library helpers nlargest and nsmallest are convenient, but for very large inputs they build the full sorted list internally, so writing the heap loop by hand is what actually keeps the memory bounded.

Quickselect partitions the array around a pivot exactly as quicksort does, but instead of recursing into both halves it continues into only the half that contains the answer. Once the pivot lands on the required index, the answer is the pivot itself.

def kth_largest(values, k):
    target = len(values) - k
    data = list(values)

    def partition(low, high):
        pivot = data[high]
        store = low
        for i in range(low, high):
            if data[i] < pivot:
                data[store], data[i] = data[i], data[store]
                store += 1
        data[store], data[high] = data[high], data[store]
        return store

    low, high = 0, len(data) - 1
    while low <= high:
        pivot_index = partition(low, high)
        if pivot_index == target:
            return data[pivot_index]
        if pivot_index < target:
            low = pivot_index + 1
        else:
            high = pivot_index - 1
    return None


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

The output is 5, which is the second largest value. The array is copied at the start so the caller's list is not reordered, which is a detail worth stating because the partition step swaps elements in place.

The mechanism is that the partition places the pivot at its final sorted position. If that position is the one you want, you are done. If it is smaller, the answer lies to the right and the left half is discarded entirely; if it is larger, the right half is discarded. The expected total work is linear because each partition is linear in the size of the interval and the intervals shrink by a constant factor on average.

The trade-off to name is the worst case. If the pivot is always the smallest or the largest remaining element, the intervals shrink by one each time and the algorithm becomes quadratic. Randomising the pivot or taking the median of three sample values makes that case astronomically unlikely. If you need a guaranteed bound, select the pivot by the median of medians, which gives linear time at the cost of a much longer implementation, or use the heap version, which has a guaranteed O of n log k bound at the price of extra memory.