Find the Kth max and min element of an array

The problem, simply

You get an array of numbers and a number k. Give back the k-th smallest number (the one that would be in position k if you sorted from small to big) and the k-th largest (position k if you sorted from big to small).

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

A story to remember it

Picture it

Examples

Way 1: the simple way

Sort, then read the two positions.

1. Sort the array.
2. The k-th smallest is at index k - 1. The k-th largest is at index n - k.

class Solution {
    public int[] kthSmallestAndLargest(int[] arr, int k) {
        int[] a = arr.clone();
        Arrays.sort(a);
        return new int[] { a[k - 1], a[a.length - k] };   // k-th smallest, k-th largest
    }
}

How long it takes: sorting is O(n log n). Extra room: O(1) besides the copy and the sort. It does much more work than needed: it puts every number in place, but we only care about two positions.

Way 2: the clever way

Use the measuring stick trick, called quickselect. It is quicksort, but after splitting it only goes into the one side that holds the position we want.

1. Let target = k - 1 (for the k-th largest, target = n - k).
2. Pick a random box as the pivot and move it to the end of the range.
3. Split: move every number not bigger than the pivot to the left side, then put the pivot right after them, at box p. The pivot is now exactly where sorting would put it.
4. If p == target, done. If target < p, keep only the left part. Otherwise keep only the right part. Repeat.

In the animation, the random pick happens to land on the last box each time, so you can follow it easily.

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

    public int[] kthSmallestAndLargest(int[] arr, int k) {
        int n = arr.length;
        int small = select(arr.clone(), k - 1);   // position k-1 in sorted order
        int large = select(arr.clone(), n - k);   // position n-k in sorted order
        return new int[] { small, large };
    }

    private int select(int[] a, int target) {
        int lo = 0, hi = a.length - 1;
        while (true) {
            int r = lo + rnd.nextInt(hi - lo + 1);   // random measuring stick
            swap(a, r, hi);
            int p = partition(a, lo, hi);            // pivot lands at its sorted place
            if (p == target) return a[p];
            if (target < p) hi = p - 1;              // answer is on the left
            else lo = p + 1;                         // answer is on the right
        }
    }

    private int partition(int[] a, int lo, int hi) {
        int pivot = a[hi];
        int i = lo - 1;                              // end of the "not bigger" side
        for (int j = lo; j < hi; j++) {
            if (a[j] <= pivot) swap(a, ++i, j);
        }
        swap(a, i + 1, hi);
        return i + 1;
    }

    private void swap(int[] a, int i, int j) {
        int t = a[i]; a[i] = a[j]; a[j] = t;
    }
}

How long it takes: on average each round throws away about half of what is left, so the work is about n + n/2 + n/4 + ..., which is less than 2n: O(n) on average. With very bad luck every pivot could be the biggest number, which is O(n^2), but random picks make that extremely unlikely. Extra room: O(1) besides the copies, since the loop does not use recursion.

Another way: a pile of size k

Keep a pile (a heap) of the k smallest numbers seen so far, with the biggest of them on top. A new number smaller than the top kicks the top out. At the end, the top is the k-th smallest. For the k-th largest, keep the k largest with the smallest on top.

This is the best choice when the numbers arrive one by one (a stream), or when k is small and n is huge, because it never needs the whole array at once.

class Solution {
    public int kthSmallest(int[] arr, int k) {
        PriorityQueue<Integer> pile = new PriorityQueue<>(Collections.reverseOrder()); // biggest on top
        for (int x : arr) {
            if (pile.size() < k) {
                pile.offer(x);
            } else if (x < pile.peek()) {
                pile.poll();                  // the top cannot be among the k smallest
                pile.offer(x);
            }
        }
        return pile.peek();
    }

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

This takes O(n log k) time and O(k) extra room.

Watch out for

• Picking the first or last box as the pivot every time is O(n^2) on an already sorted array. Pick at random.
• Quickselect shuffles the array. Work on a copy if the caller needs the original order.
• The pile for the k-th smallest is a max heap (biggest on top), not a min heap. It sounds backwards, but the top is the one that gets kicked out.

What interviewers ask next

• Can you guarantee O(n) in the worst case? Yes, with the "median of medians" pivot: split into groups of 5, take each group's middle, and pick the middle of those as the pivot. It is slower in practice, so random pivots are used.
• Many arrays of equal numbers make Lomuto partition slow. Use a three way split (smaller, equal, bigger) so a whole block of equal numbers is placed at once.
• k-th largest in a stream of numbers? The pile of size k is the answer: O(log k) per new number.
• k-th smallest in a sorted matrix, or in two sorted arrays? Those have their own tricks: binary search on the value, or a heap over the rows.
• Top k numbers, not only the k-th? Quickselect to position k, then everything left of it is the top k (in any order).

Remember this

Grown-up words: the measuring stick is the pivot, the left and right walls are the partition (Lomuto partition), the whole method is quickselect (Hoare's selection), and the pile is a binary heap used as a priority queue.