Find the k Largest Elements in an Array

Given an array of integers and a number k, return the k largest elements, largest first.

Example 1

Text

Input:  [9, 4, 7, 1, -2, 6, 5], k = 3
Output: [9, 7, 6]

Example 2

Text

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

Brute force: sort in descending order

An int[] cannot be sorted with a comparator, so the version below sorts a copy in ascending order and reads it from the end.

Java

import java.util.Arrays;

public class Main {
    static int[] kLargest(int[] nums, int k) {
        int[] sorted = nums.clone();
        Arrays.sort(sorted);
        int size = Math.min(k, sorted.length);
        int[] result = new int[size];
        for (int i = 0; i < size; i++) {
            result[i] = sorted[sorted.length - 1 - i]; // largest first
        }
        return result;
    }

    public static void main(String[] args) {
        System.out.println(Arrays.toString(kLargest(new int[]{9, 4, 7, 1, -2, 6, 5}, 3)) + " " + Arrays.toString(kLargest(new int[]{3, 2, 1, 5, 6, 4}, 2)));
    }
}

Output

[9, 7, 6] [6, 5]

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

Optimal: a min-heap of size k

This mirrors k smallest elements. Keep the k largest values seen so far in a min-heap, so its top is the smallest of them: the one to evict when a larger value arrives.

PriorityQueue is a min-heap by default, so no comparator is needed.

This version adds every value and then polls whenever the heap grows past k. poll always removes the smallest value, so the k largest remain.

Walkthrough with [9, 4, 7, 1, -2, 6, 5], k = 3 (heap contents shown sorted):

x	after add	size > 3?	poll removes	heap after
9	{9}	no		{9}
4	{4, 9}	no		{4, 9}
7	{4, 7, 9}	no		{4, 7, 9}
1	{1, 4, 7, 9}	yes	1	{4, 7, 9}
-2	{-2, 4, 7, 9}	yes	-2	{4, 7, 9}
6	{4, 6, 7, 9}	yes	4	{6, 7, 9}
5	{5, 6, 7, 9}	yes	5	{6, 7, 9}

Polling the remaining heap returns the values in ascending order, so the result array is filled from the back to put the largest first.

Java

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

public class Main {
    static int[] kLargest(int[] nums, int k) {
        PriorityQueue<Integer> heap = new PriorityQueue<>();
        for (int x : nums) {
            heap.add(x);
            if (heap.size() > k) {
                heap.poll(); // removes the smallest value
            }
        }
        int[] result = new int[heap.size()];
        for (int i = result.length - 1; i >= 0; i--) {
            result[i] = heap.poll(); // ascending order, so fill from the back
        }
        return result;
    }

    public static void main(String[] args) {
        System.out.println(Arrays.toString(kLargest(new int[]{9, 4, 7, 1, -2, 6, 5}, 3)));
        System.out.println(Arrays.toString(kLargest(new int[]{3, 2, 1, 5, 6, 4}, 2)));
        System.out.println(Arrays.toString(kLargest(new int[]{1, 2}, 0)));
    }
}

Output

[9, 7, 6]
[6, 5]
[]

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

Iterating a PriorityQueue or printing it does not give sorted order, so the code polls the values one at a time.

The kth largest element (LeetCode 215)

LeetCode 215 asks only for the kth largest value. That's the top of the same size-k min-heap, the smallest of the k largest. The method returns OptionalInt, a container that holds either one int or nothing, and it is empty when the array has fewer than k elements.

Java

import java.util.OptionalInt;
import java.util.PriorityQueue;

public class Main {
    static OptionalInt kthLargest(int[] nums, int k) {
        PriorityQueue<Integer> heap = new PriorityQueue<>();
        for (int x : nums) {
            heap.add(x);
            if (heap.size() > k) {
                heap.poll();
            }
        }
        return k > 0 && heap.size() == k ? OptionalInt.of(heap.peek()) : OptionalInt.empty();
    }

    public static void main(String[] args) {
        System.out.println(kthLargest(new int[]{3, 2, 1, 5, 6, 4}, 2));
        System.out.println(kthLargest(new int[]{3, 2, 3, 1, 2, 4, 5, 5, 6}, 4));
        System.out.println(kthLargest(new int[]{1}, 2));
    }
}

Output

OptionalInt[5]
OptionalInt[4]
OptionalInt.empty

With the whole array in memory, quickselect answers it in average O(n) time. The Java standard library has no quickselect method.

Edge cases

k = 0: empty result.
k >= n: every element, largest first.
Duplicates: kept. The 4th largest of [3, 2, 3, 1, 2, 4, 5, 5, 6] is 4, counting both 5s.