Find the k Smallest Elements in an Array

Given an array of integers and a number k, return the k smallest elements in ascending order. If k is larger than the array, return every element, sorted.

Example 1

Text

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

Example 2

Text

Input:  [3, 1], k = 5
Output: [1, 3]

Brute force: sort everything

Java

import java.util.Arrays;

public class Main {
    static int[] kSmallest(int[] nums, int k) {
        int[] sorted = nums.clone();
        Arrays.sort(sorted);
        return Arrays.copyOf(sorted, Math.min(k, sorted.length));
    }

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

Output

[-2, 1, 4] [1, 3]

Time: O(n log n). Space: O(n) for the copy. The whole array is sorted, even though only k elements are needed.

Optimal: a max-heap of size k

Walk the array while keeping the k smallest values seen so far in a max-heap. The top of that heap is the largest of the kept values: the weakest candidate, and the first to go if something smaller shows up.

For each value x:

If the heap has fewer than k elements, push x.
Otherwise, if x is smaller than the top, replace the top with x.
Otherwise, x is larger than all k kept values, so skip it.

Walkthrough with [9, 4, 7, 1, -2, 6, 5], k = 3:

x	heap before (top first)	action	heap after
9	[]	push	[9]
4	[9]	push	[9, 4]
7	[9, 4]	push	[9, 7, 4]
1	top 9	1 < 9: replace	[7, 4, 1]
-2	top 7	-2 < 7: replace	[4, 1, -2]
6	top 4	6 > 4: skip	[4, 1, -2]
5	top 4	5 > 4: skip	[4, 1, -2]

Java

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

public class Main {
    static int[] kSmallest(int[] nums, int k) {
        if (k <= 0) return new int[0];
        PriorityQueue<Integer> heap = new PriorityQueue<>(Collections.reverseOrder()); // max-heap
        for (int x : nums) {
            if (heap.size() < k) {
                heap.add(x);
            } else if (x < heap.peek()) {
                heap.poll(); // remove the current maximum
                heap.add(x);
            }
        }
        int[] result = new int[heap.size()];
        for (int i = result.length - 1; i >= 0; i--) {
            result[i] = heap.poll(); // the largest comes out first, so fill from the back
        }
        return result;
    }

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

Output

[-2, 1, 4]
[1, 3]
[1, 5]
[]

Time: O(n log k): each of the n elements costs at most one O(log k) heap operation. Space: O(k).

This is much faster than sorting when k is small compared to n. It also works when the numbers arrive as a stream too large to hold in memory, since only k values are ever stored.

PriorityQueue has no operation that replaces the top in place, so the code calls poll() and then add(x). Both are O(log k). Collections.reverseOrder() turns the min-heap into a max-heap. PriorityQueue.toString() and iteration are not sorted, so the code drains the heap with poll() to get ascending order.

Alternative: quickselect

Quickselect rearranges an array so that the element at index i is the one that would be there after sorting, everything before it is <= it, and everything after is >= it. It partitions around a random pivot like quicksort, but continues into only the side that contains index i. That runs in O(n) on average. The Java standard library has no such method, so the code below writes it. The three-way partition keeps runs of equal values from degrading it.

Java

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

public class Main {
    static final Random RANDOM = new Random(42);

    // Rearranges a so that a[0..k-1] holds the k smallest values, in any order.
    static void quickselect(int[] a, int k) {
        int lo = 0, hi = a.length - 1, target = k - 1;
        while (lo < hi) {
            int pivot = a[lo + RANDOM.nextInt(hi - lo + 1)];
            int lt = lo, i = lo, gt = hi; // a[lo..lt-1] < pivot, a[gt+1..hi] > pivot
            while (i <= gt) {
                if (a[i] < pivot) swap(a, lt++, i++);
                else if (a[i] > pivot) swap(a, i, gt--);
                else i++;
            }
            if (target < lt) hi = lt - 1;
            else if (target > gt) lo = gt + 1;
            else return; // the target index holds a value equal to the pivot
        }
    }

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

    static int[] kSmallest(int[] nums, int k) {
        int[] v = nums.clone();
        if (k <= 0) return new int[0];
        if (k >= v.length) {
            Arrays.sort(v);
            return v;
        }
        quickselect(v, k);
        int[] result = Arrays.copyOf(v, k);
        Arrays.sort(result); // sort only those k
        return result;
    }

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

Output

[-2, 1, 4]
[1, 5]

Time: O(n + k log k) on average (O(n²) in the worst case). Space: O(n) for the copy.

Choosing between them: quickselect is faster when the whole array is in memory and a copy can be modified. The heap suits streams and very small k.

Edge cases

k = 0: empty result.
k >= n: all elements, sorted.
Duplicates: kept, as in [5, 5, 5, 1] with k = 2 giving [1, 5].