Quick Select

QuickSelect finds the k-th smallest element (0-based) in expected linear time by repeatedly partitioning the array in-place, shrinking the search to the side that must contain the answer.

Context

QuickSelect is a selection algorithm derived from QuickSort that finds the k-th smallest element in an array without fully sorting it. It does this by partitioning around a pivot and then continuing only on the side that can still contain the answer.

Two important clarifications that matter in real systems:

• Indexing: In the implementation below, k is 0-based (k = 0 means “minimum”, k = n-1 means “maximum”). Mixing 0-based and 1-based “k-th” definitions is a classic bug—make it explicit in your API and checks.
• Mutation: QuickSelect is typically in-place. That’s great for memory, but it means the input array is reordered (which can be a footgun if callers assume the array stays intact).

Core idea

• Pick a pivot element.
• Partition the array so that:
  • elements ≤ pivot go to the left
  • elements > pivot go to the right
• Let j be the pivot’s final index after partitioning.
  • If j == k, you’ve found the answer: array[j].
  • If j < k, the answer must be in the right partition (search [j+1, hi]).
  • If j > k, the answer must be in the left partition (search [lo, j-1]).

After partition(array, lo, hi) returns index j, you can rely on:

• For all i in [lo, j-1]: array[i] <= array[j]
• For all i in [j+1, hi]: array[i] > array[j]

That invariant is what makes it safe to discard half the remaining range (not necessarily half the elements—just the half that can’t contain the k-th).

Implementation

class QuickSelect {
    
    /**
     * Returns the k-th smallest element (0-based) from the array.
     *
     * Contract:
     * - k must be in [0, array.length - 1]
     * - array is mutated (elements are re-ordered)
     *
     * Expected time: O(n) with decent pivot choices (e.g., median-of-three).
     * Worst-case time: O(n^2) (e.g., adversarial ordering + unlucky pivots).
     * Extra space: O(1).
    */
    public static <T extends Comparable<? super T>> int kthElement(T[] array, int k) {
        if (k < 0 || k > array.length) throw new IllegalArgumentException("k must be greater than 0 and less than array's length");
    
        var lo = 0;
        var hi = array.length - 1;
        
        while (lo < hi) {
            var j = partition(array, lo, hi);
            if (j == k)         return array[j];
            else if (j < k)     lo = j + 1;
            else                hi = j - 1;
        }
        
        return array[lo];
    }
    
    /**
     * Recursive flavor of the same algorithm.
     *
     * Tradeoff note: recursion is clean, but can overflow the stack in worst cases.
     * Iterative form is safer for production unless you can bound recursion depth.
     */
     public static <T extends Comparable<? super T>> int kthElement(T[] array, int k, int lo, int hi) {
        if (k < 0 || k > array.length) throw new IllegalArgumentException("k must be greater than 0 and less than array's length");
        
        var pivot = partition(array, lo, hi);
        
        if (pivot == k) return array[pivot];
        
        if (pivot < k) return partition(array, pivot + 1, hi);
        
        return partition(array, lo, pivot - 1);
    }
    
    private static <T extends Comparable<? super T>> int partition(T[] array, int lo, int hi) {
       var mid = lo + (hi - lo) / 2;
       var pivot = medianOfThree(array, lo, mid, hi);
       T pivotValue = array[pivot];

       swap(array, pivotIndex, hi); // move pivot to end for partitioning       
       
       var j = lo - 1;
       for (var i = lo; i < hi; i++) {
           if (array[i].compareTo(pivotValue) <= 0) {
                swap(array, ++j, i);
           }
       }
       
       swap(array, ++j, hi); // place pivot into its final position
       return j;
    }

    /**
     * Chooses the median index among lo, mid, hi to reduce the probability of worst-case pivots
     * on already-sorted or nearly-sorted data (a common real-world input shape).
     *
     * Note: This is a heuristic, not a guarantee—adversarial inputs can still force O(n^2).
     * If you need worst-case linear time, look at Median of Medians (more constant-factor cost).
     */
    private static <T extends Comparable<? super T>> int medianOfThree(T[] array, int lo, int mid, int hi) {
        T a = array[lo], b = array[mid], c = array[hi];
        
        if (a.compareTo(b) <= 0) {
            if (b.compareTo(c) <= 0) return mid; // a <= b <= c
            if (a.compareTo(c) <= 0) return hi;  // a <= c <= b
            return lo;                           // c <= a <= b
        }
        
        if (c.compareTo(b) <= 0) return mid; // c <= b <= a
        if (b.compareTo(c) <= 0) return hi; // b <= c <= a
        return lo;
    }

    private static <T extends Comparable<? super T>> void swap(T[] array, int a, int b){
        var temp = array[a];
        array[a] = array[b];
        array[b] = temp;
    }
}