Partial sorting

The nth-element operation rearranges an array so that the element at index n is the one that would be there if the whole array were sorted. Everything before it is less than or equal to it, and everything after it is greater than or equal to it, in no particular order within each side.

That is exactly what one run of quickselect leaves behind. C++ provides std::nth_element (introselect, O(n)). Java has no direct equivalent, so write it with the quickselect partition from lesson 19.02. It is the foundation for median, percentile and top-K queries that do not need sorted output.

1// after nthElement(arr, n): arr[0..n-1] <= arr[n] <= arr[n+1..]
2static void nthElement(int[] arr, int n) {
3    int lo = 0, hi = arr.length - 1;
4    while (lo < hi) {
5        int p = partition(arr, lo, hi);      // Lomuto or Hoare, random pivot recommended
6        if (p == n) return;
7        if (p < n) lo = p + 1; else hi = p - 1;
8    }
9}

Partial heap sort produces the k smallest values in sorted order. Build a min-heap of all n elements in O(n), then poll k times, each costing O(log n). The total is O(n + k log n), which beats a full sort when k is much smaller than n.

Java's PriorityQueue(Collection) constructor heapifies in O(n), which is faster than adding the elements one by one (O(n log n)).

1static int[] kSmallestSorted(int[] arr, int k) {
2    PriorityQueue<Integer> heap = new PriorityQueue<>(Arrays.stream(arr).boxed().toList());   // O(n) heapify
3    int[] out = new int[Math.min(k, arr.length)];
4    for (int i = 0; i < out.length; i++) out[i] = heap.poll();                                // O(log n) each
5    return out;
6}

A fast way to get the k smallest values in sorted order: use nth-element to place the k-th value so that the first k slots hold the k smallest values in some order, then sort just that prefix. The cost is O(n) expected for the selection plus O(k log k) for sorting the prefix.

Compare the options by size of k. If k is a tiny constant, the size-k heap is simple and works on a stream. If k is close to n, a full sort is simpler. In between, selection plus a small sort is usually fastest.

1static int[] topKSorted(int[] arr, int k) {
2    int[] copy = arr.clone();
3    nthElement(copy, k - 1);              // the k smallest now occupy copy[0..k-1]
4    Arrays.sort(copy, 0, k);              // sort only those k
5    return Arrays.copyOf(copy, k);
6}

Given points and a query point (usually the origin), return the k closest. Compare squared distances, x*x + y*y, as long, to avoid both the square root and int overflow.

Method 1: a max-heap of size k holding the best k so far. A new point is added if it is closer than the heap's root (the farthest of the current k). That costs O(n log k). Method 2: quickselect on the distance, expected O(n), when you can reorder the array and do not need the result in order.

1static int[][] kClosest(int[][] points, int k) {
2    PriorityQueue<int[]> heap = new PriorityQueue<>(
3        (a, b) -> Long.compare(dist(b), dist(a)));        // max-heap on distance
4    for (int[] p : points) {
5        heap.add(p);
6        if (heap.size() > k) heap.poll();                  // drop the farthest
7    }
8    return heap.toArray(new int[0][]);
9}
10private static long dist(int[] p) { return (long) p[0] * p[0] + (long) p[1] * p[1]; }