DSA Kth Largest Element

The problem

The Kth Largest Element is the value that would sit at position k in a descending sort. Position 1 is the maximum. Position n, for a list of n values, is the minimum.

Duplicate values count as separate entries. Two copies of 7 can hold positions 1 and 2 at once. The cases below include one duplicate case on purpose.

cases = [
    ([3, 2, 1, 5, 6, 4], 1, 6),
    ([3, 2, 1, 5, 6, 4], 2, 5),
    ([3, 2, 1, 5, 6, 4], 3, 4),
    ([3, 2, 1, 5, 6, 4], 6, 1),
    ([7, 7, 5], 2, 7),
]
for values, k, expected in cases:
    print("values", values, "k =", k, "->", expected)

Read one row at a time. Sorted downward, the first list runs 6, 5, 4, 3, 2, 1. Position 3 holds 4, and position 6 holds 1.

Intuition: tracking the largest is not enough

The first idea tracks the largest value seen so far. That idea answers k = 1 perfectly. It stores nothing else, so it has no answer ready for any larger k.

The missing values are the runners-up. A third largest hides among values the tracker already discarded. The code below shows the tracker finishing with 6 while the question asks for position 3.

values = [3, 2, 1, 5, 6, 4]
largest = values[0]
for v in values:
    if v > largest:
        largest = v
print("tracker answer (k = 1):", largest)
print("asked for k = 3, correct answer: 4")

The fix keeps more than one value. Keeping exactly k values turns out to be enough. Keeping all of them, sorted, is the brute force.

Brute force: sort everything, then pick position k

The brute force sorts the whole list downward and reads position k. Sorting is honest work, and it always gives the right answer. Its cost is ordering every value, including the many that sit far below position k.

The sort below is a selection sort, which repeatedly moves the largest remaining value to the front. A counter records its comparisons, so the cost is printed rather than asserted.

def kth_largest_brute(values, k):
    ordered = list(values)
    comparisons = 0
    for i in range(len(ordered)):
        best = i
        for j in range(i + 1, len(ordered)):
            comparisons += 1
            if ordered[j] > ordered[best]:
                best = j
        ordered[i], ordered[best] = ordered[best], ordered[i]
    return ordered[k - 1], ordered, comparisons

answer, ordered, comparisons = kth_largest_brute([3, 2, 1, 5, 6, 4], 3)
print("sorted:", ordered)
print("k = 3 answer:", answer)
print("comparisons:", comparisons)

Fifteen comparisons ordered six values so one position could be read. The values below position k were sorted for nothing. That waste is what the heap design removes.

The optimal solution: a heap holding only the k best values

The method in one sentence: keep a min-heap of the k largest values seen so far, and the heap top, its smallest member, is exactly the kth largest. A min-heap is a heap whose top value is the smallest one inside.

The heap fills with the first k values. After that, a new value enters only when it beats the heap top. Entering means replacing the top, so the heap size never grows past k.

import heapq

def kth_largest(values, k):
    heap = []
    pushes = 0
    for v in values:
        if len(heap) < k:
            heapq.heappush(heap, v)
            pushes += 1
        elif v > heap[0]:
            heapq.heapreplace(heap, v)
            pushes += 1
    return heap[0], pushes

cases = [
    ([3, 2, 1, 5, 6, 4], 1, 6),
    ([3, 2, 1, 5, 6, 4], 2, 5),
    ([3, 2, 1, 5, 6, 4], 3, 4),
    ([3, 2, 1, 5, 6, 4], 6, 1),
    ([7, 7, 5], 2, 7),
]
for values, k, expected in cases:
    answer, pushes = kth_largest(values, k)
    print("values", values, "k =", k, "->", answer, "expected:", expected, "match:", answer == expected, "pushes:", pushes)

All five cases pass. A value equal to the heap top does not enter, because the test is strictly greater. The duplicate row still works, since both copies of 7 entered while the heap was filling.

Walkthrough: k = 3 over [3, 2, 1, 5, 6, 4]

The walk visits one value at a time. The state is the heap, printed as stored. Its first entry is the top, the smallest of the current best three.

import heapq

values = [3, 2, 1, 5, 6, 4]
k = 3
heap = []
for v in values:
    if len(heap) < k:
        heapq.heappush(heap, v)
        action = "pushed"
    elif v > heap[0]:
        heapq.heapreplace(heap, v)
        action = "replaced the top"
    else:
        action = "skipped"
    print("value", v, "-", action, "- heap:", heap, "- top:", heap[0])
print("k = 3 answer:", heap[0])

The value 4 arrives last and replaces the top 3. After that swap the best three are 4, 5 and 6. The smallest of them, 4, is read straight off the top.

Complexity

The selection sort makes n times n minus 1, divided by 2, comparisons for n values, so its time is O(n^2). The heap walk makes one pass, and each push works inside a heap of size k with about log2 k levels, so its time is O(n log k).

The brute force stores a sorted copy of every value, so its extra space is O(n). The heap stores at most k values, so its extra space is O(k).

for n in (6, 12, 24):
    comparisons = n * (n - 1) // 2
    print("values", n, "- sort comparisons:", comparisons, " heap size capped at: 3")

Counts beat clocks here. At 24 values the full sort pays 276 comparisons. The heap still holds only 3 values, and each push climbs at most 2 levels inside it.

Edge cases and pitfalls

A k beyond the list length has no answer. The guard below answers None instead of reading a top that does not exist. Valid input keeps k between 1 and n.

Negative values need no special handling. Comparisons order them as usual, and the heap top logic never looks at signs. A one-value list answers that value for k = 1.

import heapq

def kth_largest(values, k):
    if k < 1 or k > len(values):
        return None
    heap = []
    for v in values:
        if len(heap) < k:
            heapq.heappush(heap, v)
        elif v > heap[0]:
            heapq.heapreplace(heap, v)
    return heap[0]

print("k = 1 is the maximum:", kth_largest([3, 2, 1, 5, 6, 4], 1))
print("k = n is the minimum:", kth_largest([3, 2, 1, 5, 6, 4], 6))
print("k too large:", kth_largest([3, 2], 5))
print("negatives:", kth_largest([-5, -1, -9], 2))
print("single value:", kth_largest([8], 1))
print("all equal:", kth_largest([4, 4, 4], 2))

The all-equal row answers 4 without a single replacement. Equal arrivals never beat the top, and the copies stored during the fill already hold the answer.