Topics › Heaps › Medium Problems

Kth smallest element in an array [use priority queue]

Medium TUF ID 2840 A2Z

Java source ↗ Edit this article ↗

Jump to code: Brute Better Optimal

Problem

Given an integer array arr and an integer k, return the k-th smallest element, counting duplicates separately (the element at index k - 1 once the array is sorted ascending). The sheet asks for a solution with a priority queue; the caller’s array must not be modified.

Constraints: 1 ≤ k ≤ n ≤ 10^5, values fit in an int, negatives and duplicates allowed.

Example 1

arr = [7, 10, 4, 3, 20, 15], k = 3
answer = 7               // sorted: 3, 4, 7, 10, 15, 20; the 3rd is 7

Example 2

arr = [5, 5, 5, 1], k = 2
answer = 5               // sorted: 1, 5, 5, 5; duplicates each take a position

Approach 1: Brute force, sort

Intuition. Sorting puts every element in its rank position, so the answer is simply the element at index k - 1. Sort a copy so the input stays untouched.

Algorithm.

1. Copy arr and sort the copy ascending.
2. Return copy[k - 1].

static int bruteForce(int[] arr, int k) {
    int[] a = arr.clone();
    Arrays.sort(a);
    return a[k - 1];
}

Complexity. Time O(n log n). Space O(n) for the copy.

We only care about the first k positions of the sorted order, and only about the last of those. The heap approaches spend effort only where it matters.

Approach 2: Better, heapify everything and pop k - 1 times

Intuition. A min-heap hands out elements in increasing order, one poll at a time. Put all elements into a min-heap, discard the k - 1 smallest, and the root is the k-th smallest. The important detail is how the heap is built: new PriorityQueue<>(collection) runs bottom-up heapify, which is O(n), whereas n separate offer calls cost O(n log n).

Algorithm.

1. Box the values into a list and build a PriorityQueue from it in one step.
2. Call poll() k - 1 times.
3. Return peek().

static int better(int[] arr, int k) {
    List<Integer> all = new ArrayList<>();
    for (int x : arr) all.add(x);
    PriorityQueue<Integer> minHeap = new PriorityQueue<>(all);    // built from a collection in O(n)
    for (int i = 1; i < k; i++) minHeap.poll();
    return minHeap.peek();
}

Complexity. Time O(n + k log n): linear heapify plus k - 1 removals of O(log n) each. Space O(n) for the heap. When k is small this is close to linear, but it always stores every element.

Approach 3: Optimal, max-heap of size k

Intuition. Scan once and keep only the k smallest values seen so far. Among those k, the one that matters is the largest: it is the gatekeeper. A new value belongs to the k smallest only if it is smaller than the gatekeeper, and then the gatekeeper is the value that drops out. A max-heap keeps its largest element at the root, so a max-heap of size k gives the gatekeeper in O(1) and replaces it in O(log k). When the scan ends, the heap holds the k smallest values of the array, and its root, the largest of them, is the k-th smallest.

This is the mirror image of the k-th largest problem, which uses a min-heap of size k. The rule to remember: the heap type is the opposite of the word in the question, because the root must be the element that is next to be kicked out.

Algorithm.

1. Create a max-heap with Collections.reverseOrder().
2. For each x: if the heap has fewer than k elements, add x. Otherwise, if x is smaller than the root, remove the root and add x. If x is not smaller, skip it, since it cannot be among the k smallest.
3. Return the root.

static int optimal(int[] arr, int k) {
    PriorityQueue<Integer> maxHeap = new PriorityQueue<>(Collections.reverseOrder());
    for (int x : arr) {
        if (maxHeap.size() < k) {
            maxHeap.offer(x);
        } else if (x < maxHeap.peek()) {             // x beats the largest of the current k smallest
            maxHeap.poll();
            maxHeap.offer(x);
        }
    }
    return maxHeap.peek();
}

Dry run on Example 1, k = 3:

The root is 7.

Complexity. Time O(n log k): at most one removal and one insertion per element on a heap of size k. Space O(k). Skipping values that are not below the root, rather than always inserting and then removing, avoids two heap operations for every element that cannot qualify, which is most of them in a long array.

Which heap approach to use? Approach 3 uses O(k) memory and works on a stream whose length is unknown, so it is the standard answer. Approach 2 is attractive when k is tiny and all data is already in memory. If no heap is required, quickselect (see the related k-th largest article) finds the answer in O(n) average time; note that the k-th smallest is the (n - k + 1)-th largest, so either solution answers both questions.

Edge cases and pitfalls

• Off by one. The k-th smallest is at sorted index k - 1; with the size-k heap the answer is the root after the scan, with no extra poll.
• Duplicates count separately. [5, 5, 5, 1] with k = 2 is 5. The strict x < peek test still works with ties: an equal value can replace the root without changing the answer, so skipping it is correct.
• Reverse comparator. Collections.reverseOrder() is safe for all int values. The popular (a, b) -> b - a overflows when the values are far apart: with Integer.MAX_VALUE and Integer.MIN_VALUE the subtraction wraps and the heap order breaks. The tests include both extremes.
• k = n. Approach 3 keeps every element and returns the maximum, Approach 2 pops n - 1 times; both are correct.
• Unboxing. peek() returns an Integer; returning it from an int method unboxes it, which is safe here because the heap is never empty when k ≥ 1.

Related problems

• K-th Largest element in an array is the mirror problem and adds quickselect.
• Kth largest element in a stream of running integers keeps the same gatekeeper heap alive across a stream.
• Sort K sorted array uses a heap of size k + 1 as a sliding window.
• Top K Frequent Elements applies the size-k heap to counts instead of values.
• Kth Smallest and Largest element in BST answers the question with an in-order traversal.
• Heaps (Theory Video) explains why building a heap from a collection is linear.