Kth Largest Element in an Array

Problem

https://leetcode.com/problems/kth-largest-element-in-an-array/

Given an integer array nums and an integer k, return the k^th largest element in the array.

Note that it is the k^th largest element in the sorted order, not the k^th distinct element.

Can you solve it without sorting?

Example 1:

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

Example 2:

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

Constraints:

• 1 <= k <= nums.length <= 10⁵

• -10⁴<= nums[i] <= 10⁴

Pattern

Heap, Sorting

Approaches

Min Heap

Explanation

The \(k`th largest element is the smallest element among the :math:`k\) largest elements. So if we maintain a min-heap containing only the \(k\) largest elements seen so far, the top of the heap is our answer.

Go through nums and push each number onto the heap until it holds \(k\) elements. After that, each new number num only matters if it is larger than the top (the smallest of the current top \(k\)). In that case, the top can no longer be one of the \(k\) largest, so we replace it with num.

Once all numbers are processed, the heap holds the \(k\) largest elements of nums, and its root heap[0] is the :math:`k`th largest.

Code

import heapq


def findKthLargest(nums: list[int], k: int) -> int:
    """Return the k-th largest element in nums."""
    heap: list[int] = []
    for num in nums:
        if len(heap) < k:
            heapq.heappush(heap, num)
        elif num > heap[0]:
            heapq.heapreplace(heap, num)
    return heap[0]

Test

>>> from kth_largest_element_in_an_array__min_heap import findKthLargest
>>> findKthLargest([3, 2, 1, 5, 6, 4], 2)
5
>>> findKthLargest([3, 2, 3, 1, 2, 4, 5, 5, 6], 4)
4

Complexity

\(n\) is the number of elements in nums

Measure	Complexity	Notes
Time	\(O(n \log k)\)	one pass through the array, pushing elements onto a heap is \(O(\log k)\)
Auxiliary Space	\(O(k)\)	min heap

kth_largest_element_in_an_array__min_heap.findKthLargest(nums: list[int], k: int) → int

Return the k-th largest element in nums.