Kth Largest Element in an Array

medium heap sorting

Problem

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

Inputnums = [9, 4, 1, 7, 6, 5], k = 3

Output6

The three largest values are {9, 7, 6}; the smallest of those is 6.

nums (comma-separated)

import heapq

def find_kth_largest(nums, k):
    heap = []
    for x in nums:
        heapq.heappush(heap, x)
        if len(heap) > k:
            heapq.heappop(heap)
    return heap[0]

function findKthLargest(nums, k) {
  const heap = new MinHeap();
  for (const x of nums) {
    heap.push(x);
    if (heap.size() > k) heap.pop();
  }
  return heap.peek();
}

class Solution {
    public int findKthLargest(int[] nums, int k) {
        PriorityQueue<Integer> heap = new PriorityQueue<>();
        for (int x : nums) {
            heap.offer(x);
            if (heap.size() > k) heap.poll();
        }
        return heap.peek();
    }
}

int findKthLargest(vector<int>& nums, int k) {
    priority_queue<int, vector<int>, greater<int>> heap;
    for (int x : nums) {
        heap.push(x);
        if ((int)heap.size() > k) heap.pop();
    }
    return heap.top();
}

Explanation

We want the kth largest number. We could sort the whole array, but we do not actually need the full order — we only care about the top k values.

So we keep a min-heap that holds at most k numbers. A min-heap is a structure that always lets us peek at and remove its smallest item quickly.

We push numbers in one by one. Whenever the heap grows past size k, we pop the smallest. That throws away numbers that are too small to be in the top k.

After processing everything, the heap holds exactly the k largest values, and the smallest of those — the heap's top — is the answer.

Example: nums = [9,4,1,7,6,5], k = 3. The heap settles on the three biggest, {5,6,7,9} trimmed to {6,7,9}, and its smallest, 6, is the 3rd largest.

Time: O(n log k) Space: O(k)