Kth Largest Element in an Array (LeetCode 215) Solution

The problem

You get an integer array nums and an integer k. Return the kth largest element in sorted order (not the kth distinct one: duplicates count).

Input:  nums = [3, 2, 1, 5, 6, 4], k = 2
Output: 5

Sorted from largest, the array is 6, 5, 4, 3, 2, 1, and the second one is 5.

How to think about it

Start with brute force. Sort the array and return nums[-k]. That’s O(n log n) and perfectly correct. Interviewers accept it as a first answer, then ask you to do better, because you’re sorting everything to learn about k items.

The key insight: keep only the k largest seen so far. Scan the array once with a min-heap capped at size k. Push each number; if the heap grows past k, pop its smallest. Anything popped is smaller than k other numbers, so it can’t be the kth largest. After the scan, the heap holds exactly the k largest values, and the root (the smallest of them) is the kth largest.

Why a min-heap, not a max-heap? The thing you throw away is always the smallest of your current candidates, and a min-heap gives you that in O(log k). A max-heap of all n numbers works too (pop k times), but it costs O(n) space. This inversion, min-heap for the k largest, is the lesson’s signature trick.

Which pattern is this? Top-K with a heap. The tell is “kth largest”, “top k” or “k closest”. If an interviewer pushes for average O(n), mention Quickselect (partition like quicksort, recurse into one side only).

Approach

Start with an empty heap.
For each number, push it.
If the heap has more than k items, pop the smallest.
After the loop, return h[0].

Solution

import heapq
def find_kth_largest(nums, k):
    h = []
    for x in nums:
        heapq.heappush(h, x)
        if len(h) > k: heapq.heappop(h)                      # drop the smallest; it can't be in the top k
    return h[0]                                              # smallest of the k largest = kth largest
# find_kth_largest([3,2,1,5,6,4], 2) -> 5

Dry run

On [3, 2, 1, 5, 6, 4] with k = 2:

`x`	Popped	Heap after (sorted)
3	—	[3]
2	—	[2, 3]
1	1	[2, 3]
5	2	[3, 5]
6	3	[5, 6]
4	4	[5, 6]

The heap ends as the two largest, 5 and 6. Its root is 5.

Complexity

Time: O(n log k). Each of the n numbers does a push and maybe a pop on a heap of at most k + 1 items. When k is small compared with n, this is much better than sorting.
Space: O(k). The heap never holds more than k + 1 items.

Edge cases and pitfalls

Duplicates count. [3,2,3,1,2,4,5,5,6] with k = 4 gives 4: the order is 6, 5, 5, 4.
k = 1 gives the maximum; k = n gives the minimum.
Single element with k = 1 returns that element.
Using a max-heap with size k is the most common mix-up: it keeps the k smallest, and the answer comes out wrong.
Quickselect worst case is O(n²) with bad pivots. Use a random pivot if you write it.

Kth Largest Element in a Stream (LeetCode 703): the same heap, kept alive as numbers arrive.
K Closest Points to Origin (LeetCode 973): top-K again, keyed by distance.
Top K Frequent Elements (LeetCode 347): top-K keyed by frequency.