LeetCode 215 – Kth Largest Element in an Array

Kth Largest Element is a menu problem: there are three standard answers with genuinely different trade-offs, and the interview is about choosing among them out loud rather than producing the cleverest one. Quickselect is the O(n) answer, and it is also the one with a worst case worth being honest about.

The problem

Return the kth largest element in an array — in sorted order, not the kth distinct value.

[3,2,1,5,6,4],       k = 2  -> 5
[3,2,3,1,2,4,5,5,6], k = 4  -> 4      duplicates COUNT
[1],                 k = 1  -> 1
[2,1],               k = 2  -> 1

The second example is the one to confirm. Sorted descending it is 6,5,5,4,3,2,2,1,1, and the 4th is 4 — the two 5s occupy positions 2 and 3. If the question wanted distinct values the answer would be 3, and that is a different problem.

The three answers

Approach	Time	Space	When
Sort, take `nums[n-k]`	`O(n log n)`	`O(1)`	one query, clarity matters
Min-heap of size `k`	`O(n log k)`	`O(k)`	streaming, or `k` ≪ `n`
Quickselect	`O(n)` average	`O(1)`	one query, `n` large

Say all three and pick one. Sorting is a perfectly respectable answer and the log n factor is often irrelevant; the heap is the right answer when the data arrives as a stream and cannot be sorted; quickselect is the one the problem is fishing for.

The heap: keep the k largest, evict the rest

Maintain a min-heap of size k. Its root is the smallest of the k largest seen so far — so when a bigger element arrives, the root is the one to discard.

push each element
if size > k: pop the smallest

at the end the heap holds the k largest,
and its root is the kth largest

A min-heap for a "largest" question is the counter-intuitive part and worth stating: you want cheap access to the weakest survivor, because that is the one being replaced.

Java

class Solution {
    public int findKthLargest(int[] nums, int k) {
        // A MIN-heap: its root is the weakest of the k largest, so it is what gets evicted.
        PriorityQueue<Integer> smallestOfTheBest = new PriorityQueue<>();

        for (int num : nums) {
            smallestOfTheBest.offer(num);
            if (smallestOfTheBest.size() > k) smallestOfTheBest.poll();
        }

        return smallestOfTheBest.peek();
    }
}

Pushing then evicting, rather than comparing before pushing, keeps it to two lines and is the same number of heap operations. The heap never exceeds k + 1 elements.

This is the version to write when k is small or the input is a stream — it needs only one pass and never holds more than k values, so it works on data that does not fit in memory.

Quickselect

Quicksort partitions around a pivot and then sorts both halves. Quickselect partitions and then recurses into only the half that contains the answer — which is what turns n log n into n.

the kth largest sits at index n - k in ascending order

partition -> the pivot lands at its final index p
  p == target   -> done
  p < target    -> recurse right
  p > target    -> recurse left

Each step discards half the array on average, so the work is n + n/2 + n/4 + … = 2n. That geometric sum is the whole complexity argument and it is short enough to give.

    int findKthLargestQuickselect(int[] nums, int k) {
        int target = nums.length - k;         // index in ASCENDING order
        int lo = 0, hi = nums.length - 1;
        Random rng = new Random();

        while (true) {
            // Random pivot: without it, a sorted input is the O(n^2) worst case.
            int pivotIndex = lo + rng.nextInt(hi - lo + 1);
            int p = partition(nums, lo, hi, pivotIndex);

            if (p == target) return nums[p];
            if (p < target) lo = p + 1;
            else hi = p - 1;
        }
    }

    private int partition(int[] nums, int lo, int hi, int pivotIndex) {
        int pivot = nums[pivotIndex];
        swap(nums, pivotIndex, hi);           // park the pivot at the end

        int write = lo;
        for (int i = lo; i < hi; i++) {
            if (nums[i] < pivot) swap(nums, write++, i);
        }

        swap(nums, write, hi);                // put the pivot in its final place
        return write;
    }

    private void swap(int[] nums, int i, int j) {
        int tmp = nums[i]; nums[i] = nums[j]; nums[j] = tmp;
    }

The random pivot is not optional. Always taking the last element makes an already sorted array partition into pieces of size n−1 and 0 every time, which is O(n²) — and "already sorted" is the most likely shape of real input. Randomising makes that worst case require an adversary who knows your seed.

Note that quickselect reorders the caller's array. That is a side effect the problem tolerates and a real API would not; say so rather than leaving it silent.

Python

import heapq


class Solution:
    def findKthLargest(self, nums: list[int], k: int) -> int:
        # nlargest is a size-k heap under the hood: O(n log k), one pass.
        return heapq.nlargest(k, nums)[-1]

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

        for num in nums:
            heapq.heappush(heap, num)
            if len(heap) > k:
                heapq.heappop(heap)           # evict the weakest survivor

        return heap[0]

heapq is a min-heap only, which happens to be exactly what this needs. For a max-heap the idiom is to push negated values — worth knowing, and not needed here.

heapq.nlargest(k, nums)[-1] is the one-liner and does the same work. Use it, but be ready to write the loop: an interviewer asking this wants to know you could implement the eviction rule.

Complexity

Approach	Average	Worst	Space
Sort	`O(n log n)`	`O(n log n)`	`O(1)` to `O(n)`
Heap of size k	`O(n log k)`	`O(n log k)`	`O(k)`
Quickselect	`O(n)`	`O(n²)`	`O(1)`

Quickselect's worst case is real and the randomised pivot makes it vanishingly unlikely rather than impossible. There is a deterministic O(n) worst-case algorithm — median-of-medians — and its constant factor is bad enough that it is almost never used in practice. Naming it and explaining why it stays on the shelf is a better answer than pretending quickselect is worst-case linear.

The pattern

"Keep the best k with a heap of the opposite polarity" is the reusable move. It is what Top K Frequent Elements (347) does, and the same reasoning behind Kth Smallest in a Sorted Matrix (378) and Find Median from Data Stream (295), which balances two heaps of opposite direction.

Quickselect generalises less often but is worth recognising: any "find the kth something" over an unsorted array where you may reorder it.

What the interviewer is checking

That kth largest counts duplicates, not distinct values.
That you name all three approaches and choose deliberately.
Why a min-heap answers a "largest" question.
That quickselect recurses into one side only, and the 2n argument.
The random pivot, and what a sorted input does without it.
That quickselect mutates the caller's array.
k = 1, k = n, and a single-element array.