k largest element in an array

The problem, simply

You get an array of n numbers and a number k. Give back the k biggest numbers, from biggest to smallest.

For [1, 23, 12, 9, 30, 2, 50] and k = 3, the answer is [50, 30, 23].

A story to remember it

Picture it

Examples

Way 1: the simple way

Sort all the numbers from biggest to smallest, then take the first k.

1. Copy the numbers into Integer boxes so Java can sort them backwards.
2. Sort with Collections.reverseOrder().
3. Return the first k.

import java.util.Arrays;
import java.util.Collections;

class Solution {
    public int[] kLargest(int[] arr, int k) {
        Integer[] boxed = new Integer[arr.length];
        for (int i = 0; i < arr.length; i++) boxed[i] = arr[i];
        Arrays.sort(boxed, Collections.reverseOrder());   // biggest first
        int[] out = new int[k];
        for (int i = 0; i < k; i++) out[i] = boxed[i];
        return out;
    }
}

How long it takes: sorting is O(n log n). Extra room: the boxed copy, O(n). We sorted all n numbers but only needed k of them.

Way 2: the clever way

Do what the club does. Keep a min heap that never holds more than k numbers. Its top is the smallest of the k biggest so far, the one by the door.

1. For each number x:
   • If the heap has fewer than k numbers, add x.
   • Else if x is bigger than the top, remove the top and add x.
   • Else skip x. It cannot be in the top k.
2. At the end the heap holds the answer. Poll them out (smallest first) and fill the answer from the back.

import java.util.PriorityQueue;

class Solution {
    public int[] kLargest(int[] arr, int k) {
        PriorityQueue<Integer> club = new PriorityQueue<>();   // min heap
        for (int x : arr) {
            if (club.size() < k) {
                club.offer(x);                   // a free chair
            } else if (x > club.peek()) {
                club.poll();                     // the weakest member leaves
                club.offer(x);
            }                                    // else: x walks on
        }
        int[] out = new int[k];
        for (int i = k - 1; i >= 0; i--) out[i] = club.poll(); // smallest goes last
        return out;
    }
}

How long it takes: each of the n numbers costs at most one poll and one offer on a heap of size k, so O(n log k). Extra room: the heap, O(k). When k is small this is much faster than sorting, and it works on a stream of numbers that never ends.

Another way: quickselect

Quickselect is quicksort that only follows one side. Pick a pivot, move smaller numbers left and bigger numbers right. If the pivot lands at index n - k, everything to its right is the top k. If not, repeat only on the side that holds index n - k.

import java.util.Arrays;
import java.util.Random;

class Solution {
    private final Random rand = new Random();

    public int[] kLargest(int[] arr, int k) {
        int[] a = arr.clone();
        int n = a.length, target = n - k, lo = 0, hi = n - 1;
        while (lo < hi) {
            int p = partition(a, lo, hi);
            if (p == target) break;
            if (p < target) lo = p + 1; else hi = p - 1;
        }
        int[] out = Arrays.copyOfRange(a, target, n);
        Arrays.sort(out);                        // only k numbers to sort
        for (int i = 0, j = k - 1; i < j; i++, j--) { int t = out[i]; out[i] = out[j]; out[j] = t; }
        return out;
    }

    private int partition(int[] a, int lo, int hi) {
        int r = lo + rand.nextInt(hi - lo + 1);  // random pivot avoids bad cases
        int t = a[r]; a[r] = a[hi]; a[hi] = t;
        int pivot = a[hi], i = lo;
        for (int j = lo; j < hi; j++) {
            if (a[j] < pivot) { t = a[i]; a[i] = a[j]; a[j] = t; i++; }
        }
        t = a[i]; a[i] = a[hi]; a[hi] = t;
        return i;
    }
}

This is O(n) on average plus O(k log k) to order the answer, but O(n^2) in the unlucky worst case, and it needs all the numbers in memory at once. The heap is the safer answer when numbers arrive one at a time.

Watch out for

• Polling the min heap gives the answer smallest first. Fill the output from the back, or reverse it.
• k larger than n: decide what to return. The loop above would leave empty slots, so check k <= n first.
• Putting all n numbers into a max heap and polling k times is correct, but with one by one offer it is O(n log n) to build, not O(n).
• Do not compare Integer objects with ==. Use >, <, or equals.

What interviewers ask next

• What if the numbers come in a stream and never stop? The min heap of size k still works. It never grows past k.
• Only the k-th largest, not all of them? Same heap, return peek() at the end.
• Which is better, heap or quickselect? Quickselect is faster on average for one fixed array. The heap has no bad worst case, uses only O(k) room, and handles streams.
• K most frequent words or numbers? Count with a hash map, then run the same min heap of size k on the counts.
• K closest points to the origin? A max heap of size k on distance, because you kick out the farthest.

Remember this

Grown-up words: the club is a min heap (priority queue) of size k, the kid by the door is the heap top, kicking out and sitting down are poll and offer, and following one side of the split is quickselect.