LeetCode 215: Kth Largest Element in an Array

Problem Restatement

We are given an integer array nums and an integer k.

We need to return the kth largest element in the array.

Important detail: this means the kth largest element in sorted order, not the kth distinct element. Duplicates count separately. The official constraints are 1 <= k <= nums.length <= 10^5 and -10^4 <= nums[i] <= 10^4.

Input and Output

Example function shape:

def findKthLargest(nums: list[int], k: int) -> int:
    ...

Examples

Example 1:

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

Sorted descending:

[6, 5, 4, 3, 2, 1]

The second largest element is:

5

Example 2:

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

Sorted descending:

[6, 5, 5, 4, 3, 3, 2, 2, 1]

The fourth largest element is:

4

Notice that both 5s count separately.

First Thought: Sort the Array

The simplest solution is to sort the array in descending order and return index k - 1.

class Solution:
    def findKthLargest(self, nums: list[int], k: int) -> int:
        nums.sort(reverse=True)
        return nums[k - 1]

This is correct and short.

Problem With Sorting

Sorting gives the complete order of all elements.

But we only need one element.

For n elements, sorting takes:

O(n log n)

That is acceptable in many cases, but the problem asks whether we can solve it without fully sorting.

Two common alternatives are:

Heap Idea

Keep only the largest k elements seen so far.

Use a min-heap.

The heap stores the current top k largest values.

The smallest value in this heap is the kth largest among the values seen so far.

For each number:

1. Push it into the heap.
2. If the heap size becomes larger than k, remove the smallest value.
3. At the end, the heap root is the answer.

Heap Walkthrough

Use:

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

Process values:

At the end, the heap contains the two largest numbers:

[5, 6]

The smallest among them is:

5

So the second largest element is 5.

Heap Correctness

The heap is maintained so that it contains the largest k values among all elements processed so far.

When a new value is inserted and the heap grows beyond size k, removing the smallest value preserves the largest k values.

After all numbers are processed, the heap contains the largest k values in the entire array. The smallest value among these k values is exactly the kth largest element.

Heap Complexity

Heap Implementation

import heapq

class Solution:
    def findKthLargest(self, nums: list[int], k: int) -> int:
        heap = []

        for num in nums:
            heapq.heappush(heap, num)

            if len(heap) > k:
                heapq.heappop(heap)

        return heap[0]

Quickselect Idea

Quickselect is based on partitioning.

Instead of sorting the full array, we only search the side that contains the answer.

If the array were sorted ascending, the kth largest element would be at index:

target = len(nums) - k

Example:

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

Sorted ascending:

[1, 2, 3, 4, 5, 6]

The second largest is at index:

6 - 2 = 4

The value at index 4 is 5.

Partitioning

Pick a pivot.

Rearrange the array so that:

After partitioning, if the pivot lands at target, we found the answer.

If the pivot lands before target, search the right side.

If the pivot lands after target, search the left side.

Quickselect Algorithm

1. Convert kth largest to ascending index: target = len(nums) - k.
2. Set search range [left, right].
3. Partition the range.
4. Compare pivot index with target.
5. Narrow the search range.
6. Return nums[target].

Quickselect Implementation

import random

class Solution:
    def findKthLargest(self, nums: list[int], k: int) -> int:
        target = len(nums) - k
        left = 0
        right = len(nums) - 1

        while left <= right:
            pivot_index = random.randint(left, right)
            pivot_index = self.partition(nums, left, right, pivot_index)

            if pivot_index == target:
                return nums[pivot_index]

            if pivot_index < target:
                left = pivot_index + 1
            else:
                right = pivot_index - 1

        return -1

    def partition(self, nums: list[int], left: int, right: int, pivot_index: int) -> int:
        pivot_value = nums[pivot_index]

        nums[pivot_index], nums[right] = nums[right], nums[pivot_index]

        store_index = left

        for i in range(left, right):
            if nums[i] < pivot_value:
                nums[store_index], nums[i] = nums[i], nums[store_index]
                store_index += 1

        nums[store_index], nums[right] = nums[right], nums[store_index]

        return store_index

Code Explanation

Convert kth largest into sorted ascending index:

target = len(nums) - k

Choose a random pivot to reduce the chance of bad partitions:

pivot_index = random.randint(left, right)

Partition the current range:

pivot_index = self.partition(nums, left, right, pivot_index)

If the pivot lands exactly where we need:

return nums[pivot_index]

Otherwise, search only one side:

if pivot_index < target:
    left = pivot_index + 1
else:
    right = pivot_index - 1

Inside partition, move the pivot to the end:

nums[pivot_index], nums[right] = nums[right], nums[pivot_index]

Place smaller elements on the left:

if nums[i] < pivot_value:
    nums[store_index], nums[i] = nums[i], nums[store_index]
    store_index += 1

Move the pivot into its final position:

nums[store_index], nums[right] = nums[right], nums[store_index]

Return that final pivot index.

Quickselect Correctness

After each partition, the pivot is placed at the same index it would occupy in a fully sorted ascending array.

All elements to its left are smaller than the pivot.

All elements to its right are greater than or equal to the pivot.

If the pivot index equals target, the pivot is the kth largest element.

If the pivot index is smaller than target, the target element must be on the right.

If the pivot index is larger than target, the target element must be on the left.

The algorithm discards the side that cannot contain the answer, while preserving the side that must contain it. Therefore it eventually returns the correct element.

Quickselect Complexity

Random pivot selection makes the bad case unlikely in practice.

Testing

def run_tests():
    s = Solution()

    assert s.findKthLargest([3,2,1,5,6,4], 2) == 5
    assert s.findKthLargest([3,2,3,1,2,4,5,5,6], 4) == 4
    assert s.findKthLargest([1], 1) == 1
    assert s.findKthLargest([2,1], 1) == 2
    assert s.findKthLargest([2,1], 2) == 1
    assert s.findKthLargest([-1,-1], 2) == -1

    print("all tests passed")

run_tests()