LeetCode Kth Largest Element in an Array Solution

Kth Largest Element in an Array (LeetCode 215) is one of the most common interview questions on arrays and heaps. The sort-based answer takes one line, but the interviewer almost always follows up with "can you do better?", so it is worth knowing the heap solution too.

Problem

Given an integer array nums and an integer k, return the kth largest element. This means the kth largest in sorted order, not the kth distinct value, so duplicates count separately. You can assume 1 ≤ k ≤ length of the array.

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

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

In the second example the array sorted from largest to smallest is 6, 5, 5, 4, ... so the 4th largest is 4. The two 5s are counted separately.

Approach 1: sort

Sort the array in ascending order. The largest element is at index n − 1, the second largest at n − 2, and so the kth largest is at index n - k. It is simple and correct, and takes O(n log n) time.

Approach 2: a min-heap of size k

We do not need the whole array sorted, only the top k values. Keep a min-heap that holds the k largest numbers seen so far. For each number, add it to the heap, and if the heap now has more than k elements, remove the smallest. At the end the heap holds the k largest values, and the smallest of those, at the top of the heap, is the answer.

Trace it on [3,2,1,5,6,4] with k = 2. After 3 and 2 the heap is {2, 3}. Adding 1 makes three items, so 1 is removed. Adding 5 removes 2, leaving {3, 5}. Adding 6 removes 3, leaving {5, 6}. Adding 4 removes 4 itself. The top of the heap is 5.

Solution in Java

import java.util.Arrays;
import java.util.PriorityQueue;

class Solution {
    // Approach 1: sort, O(n log n)
    public int findKthLargestSort(int[] nums, int k) {
        Arrays.sort(nums);
        return nums[nums.length - k];
    }

    // Approach 2: min-heap of size k, O(n log k)
    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();   // drop the smallest
        }
        return heap.peek();
    }
}

Complexity

• Sorting: O(n log n) time. Java sorts primitive arrays in place, so the extra space is small.
• Min-heap: O(n log k) time and O(k) space. When k is much smaller than n, this is clearly faster, and it also works when the numbers arrive as a stream and you cannot store them all.

Going further: Quickselect

Quickselect uses the partition step from quick sort, but only recurses into the side that contains the answer. Its average time is O(n), though a bad pivot can make it O(n²). Picking the pivot at random makes that very unlikely. If the interviewer asks for linear time on average, this is the answer they are looking for.

For more on heaps, try the heaps and hashing MCQs.

Tip: try to solve the problem yourself before reading the solution. Use the code as a reference to check your approach, and make sure you understand why it works.

More LeetCode solutions All coding problems MCQ practice