See all coding puzzles and patterns here.

Problem Definition

Given a list of integers, return the $k^{th}$ largest element in the array.


Examples

Some test cases to verify our solution.

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

Solution

We use a min heap to track the $k$ largest elements. We iterate over the entire array to look at all elements. Then we simply return the lowest element in the min heap.

def kth_largest(nums: list[int], k: int) -> int:
    min_heap = nums[:k]
    heapq.heapify(min_heap)
    for n in nums[k:]:
        if n > min_heap[0]:
            heapq.heappop(min_heap)
            heapq.heappush(min_heap, n)
    return min_heap[0]

Complexity

Let $n =$ len(nums).

Then runtime complexity should be $\Theta(k + (n - k) \cdot log(k)) = \Theta(k + n \cdot log(k) - k \cdot log(k))$.

Auxiliary space complexity should be $\Theta(k)$ since we're making a copy of the first $k$ items as the min-heap.


Destructive Unoptimized Solution

We can convert the entire input array into a min heap and then pop off the first $n-k$ elements to get the $k^{th}$ largest element. This implementation is also destructive of the input array.

def kth_largest(nums: list[int], k: int) -> int:
    heapq.heapify(nums)
    for _ in range(len(nums) - k):
        heapq.heappop(nums)
    return nums[0]

Destructive Unoptimized Solution Complexity

Let $n = $ len(nums).

Runtime complexity should be $\Theta(n + (n - k) \cdot log(n)) = \Theta(n + n \cdot log(n) - k \cdot log(n))$.

Auxiliary space complexity should be $\Theta(1)$ since we are consuming the input array.


Solution

Alternatively, we can create a max heap and pop the first $k$ elements to get the $k^{th}$ largest element.

While it's possible to just invert each value during heap constructions followed by inverting the value again prior to returning, we can also use the hidden _heapify_max and _heappop_max functions in the heapq library to implement a max heap.

def kth_largest_max_heap(nums: list[int], k: int) -> int:
    heapq._heapify_max(nums)
    for _ in range(k - 1):
        heapq._heappop_max(nums)
    return nums[0]

Destructive Solution Complexity

Let $n = $ len(nums).

Runtime complexity should be $\Theta(n + k \cdot log(n))$.

Auxiliary space complexity should be $\Theta(1)$ since we're consuming the input array.