Kth smallest and largest element in an unsorted array

The problem, simply

You get an unsorted array and a number k. Find two numbers:

• the k-th smallest: the number that would be at position k if you sorted from small to big,
• the k-th largest: the number that would be at position k if you sorted from big to small.

For [7, 10, 4, 3, 20, 15] and k = 3, the sorted order is [3, 4, 7, 10, 15, 20]. The 3rd smallest is 7 and the 3rd largest is 10.

A story to remember it

Picture it

Examples

Way 1: the simple way

Sort a copy, then read from both ends.

1. Copy the array and sort it.
2. The k-th smallest is at index k - 1.
3. The k-th largest is at index n - k.

import java.util.Arrays;

class Solution {
    public int[] kthSmallestAndLargest(int[] arr, int k) {
        int n = arr.length;
        if (k < 1 || k > n) throw new IllegalArgumentException("bad k");
        int[] a = arr.clone();               // do not change the caller's array
        Arrays.sort(a);
        return new int[] { a[k - 1], a[n - k] };
    }
}

How long it takes: sorting is O(n log n). Extra room: the copy, O(n).

Way 2: the clever way

For the k-th smallest, keep the k smallest numbers seen so far in a max heap. Its top is the biggest of those k, the one that gets pushed off when a smaller number shows up.

1. For each number x:
   • If the heap has fewer than k numbers, add x.
   • Else if x is smaller than the top, remove the top and add x.
   • Else skip x.
2. The top of the heap is the k-th smallest.

For the k-th largest, flip everything: a min heap of the k largest, push off the top when a bigger number comes, and read the top.

The same run with a min heap for the 3rd largest ends with the heap holding 10, 15 and 20, and its top, 10, is the answer.

import java.util.Collections;
import java.util.PriorityQueue;

class Solution {
    public int kthSmallest(int[] arr, int k) {
        PriorityQueue<Integer> heap = new PriorityQueue<>(Collections.reverseOrder()); // max heap
        for (int x : arr) {
            if (heap.size() < k) heap.offer(x);
            else if (x < heap.peek()) { heap.poll(); heap.offer(x); }
        }
        return heap.peek();
    }

    public int kthLargest(int[] arr, int k) {
        PriorityQueue<Integer> heap = new PriorityQueue<>();  // min heap
        for (int x : arr) {
            if (heap.size() < k) heap.offer(x);
            else if (x > heap.peek()) { heap.poll(); heap.offer(x); }
        }
        return heap.peek();
    }
}

How long it takes: each number does at most one poll and one offer on a heap of size k, so O(n log k) for each answer. Extra room: the heap, O(k).

Another way: quickselect

Pick a random pivot. Put smaller numbers to its left and bigger ones to its right. Now the pivot sits exactly where it would be in the sorted array. If that spot is the one you want, stop. Otherwise keep going only in the half that holds your spot. The k-th largest is the (n - k + 1)-th smallest, so one function serves both.

import java.util.Random;

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

    // the number that would sit at index target (0 based) after sorting
    public int select(int[] arr, int target) {
        int[] a = arr.clone();
        int lo = 0, hi = a.length - 1;
        while (true) {
            int p = partition(a, lo, hi);
            if (p == target) return a[p];
            if (p < target) lo = p + 1; else hi = p - 1;
        }
    }

    public int[] kthSmallestAndLargest(int[] arr, int k) {
        return new int[] { select(arr, k - 1), select(arr, arr.length - k) };
    }

    private int partition(int[] a, int lo, int hi) {
        int r = lo + rand.nextInt(hi - lo + 1);   // random pivot
        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 and O(n^2) in a very unlucky run. The random pivot makes the bad case almost impossible.

Watch out for

• Off by one: the k-th smallest is index k - 1 and the k-th largest is index n - k.
• Check 1 <= k <= n before you start.
• Sorting or quickselect in place changes the caller's array. Clone it if that matters.
• Quickselect with the last item as a fixed pivot is O(n^2) on an already sorted array. Pick the pivot at random.

What interviewers ask next

• Can you guarantee O(n) in the worst case? Yes, with median of medians: split into groups of 5, take each group's median, and pick the median of those as the pivot. It always throws away a good share of the numbers, but the constant is large, so random quickselect wins in practice.
• What if the numbers arrive in a stream? Quickselect needs all of them. The heap of size k works on a stream.
• What if you get many different k questions on the same array? Sort once in O(n log n), then every answer is O(1).
• What if the array is a sorted matrix or k sorted lists? Use a min heap that holds one candidate per row or list, and pop k times.

Remember this

Grown-up words: the bench is a heap (priority queue) of size k, the end seat is the heap top, the third shortest kid is the k-th order statistic, and the pivot game is quickselect.