Topics › Heaps › Medium Problems

Sort K sorted array

Easy TUF ID 2409 A2Z

Java source ↗ Edit this article ↗

Jump to code: Brute Better Optimal

Problem

You are given an array arr of n integers and an integer k. The array is k-sorted: every element is at most k positions away from the position it would occupy in sorted order. Sort the array in place, using the guarantee to beat a general-purpose sort.

Constraints: 0 ≤ n ≤ 10^5, 0 ≤ k < n in the usual statement (the code also accepts k ≥ n), values fit in an int, duplicates allowed. Such arrays appear in practice when data is “almost in order”, for example timestamps that arrive slightly out of sequence.

Example 1

arr = [6, 5, 3, 2, 8, 10, 9], k = 3
answer = [2, 3, 5, 6, 8, 9, 10]      // 6 moves from index 0 to 3, 2 from 3 to 0: both exactly 3 away

Example 2

arr = [10, 9, 8, 7, 4, 70, 60, 50], k = 4
answer = [4, 7, 8, 9, 10, 50, 60, 70]

Approach 1: Brute force, ignore k

Intuition. Any sorting algorithm sorts a k-sorted array. It is the baseline: correct, simple and O(n log n), but it wastes the information that nothing is far from home.

Algorithm.

1. Call Arrays.sort(arr).

static void bruteForce(int[] arr, int k) {
    Arrays.sort(arr);
}

Complexity. Time O(n log n). Space O(log n) for the library sort’s recursion.

Approach 2: Better, insertion sort

Intuition. Insertion sort takes each element and shifts it left past the larger elements before it. Its cost is the number of such shifts, which is small when the array is nearly sorted. How small here? Take an element x at index i whose sorted position is p, so p ≥ i - k. Every larger element y that sits before x must end up after position p. If there are c of them, they occupy c different final positions after p, so the last of them finishes at p + c or beyond. It started before index i and moved at most k, so p + c ≤ (i - 1) + k ≤ p + 2k - 1, which gives c < 2k. Each element shifts fewer than 2k times, so the whole sort is O(nk).

Algorithm.

1. For i from 1 to n - 1: save x = arr[i].
2. Walk j left from i - 1 while arr[j] > x, copying arr[j] one step right.
3. Write x into arr[j + 1].

static void better(int[] arr, int k) {
    for (int i = 1; i < arr.length; i++) {
        int x = arr[i], j = i - 1;
        while (j >= 0 && arr[j] > x) {
            arr[j + 1] = arr[j];                     // shift the larger element one step right
            j--;
        }
        arr[j + 1] = x;
    }
}

Complexity. Time O(nk). Space O(1). This beats sorting when k is smaller than log n, and it is the method to use for k = 1 or 2. For larger k the cost grows linearly in k; we want a logarithmic dependence instead.

Approach 3: Optimal, min-heap window of size k + 1

Intuition. Ask which element belongs at position 0. It is the overall minimum, and since it is at most k away from position 0, it is somewhere among arr[0..k]. So the minimum of the first k + 1 elements is the answer for position 0, and nothing beyond index k needs to be looked at yet.

The same holds for every position w: the element that belongs there sits at an index no later than w + k. Slide a window along the array with a min-heap. After reading index w + k, the heap contains every element of arr[0..w+k] that has not been placed yet. That set contains the true occupant of position w, and the true occupant is the smallest element not yet placed, so polling the heap returns it. The heap never holds more than k + 1 elements, so each operation costs O(log k).

The output can go straight back into arr: when we write position write, the reader is already at write + k, so the slot being overwritten has already been copied into the heap.

Algorithm.

1. Create an empty min-heap and set write = 0.
2. For read from 0 to n - 1: add arr[read]. If the heap now has more than k elements, poll the minimum into arr[write++].
3. When the input is exhausted, poll the rest of the heap into arr[write++] in order.

static void optimal(int[] arr, int k) {
    PriorityQueue<Integer> heap = new PriorityQueue<>();
    int write = 0;
    for (int read = 0; read < arr.length; read++) {
        heap.offer(arr[read]);
        if (heap.size() > k) arr[write++] = heap.poll();   // write <= read, so that slot was already read
    }
    while (!heap.isEmpty()) arr[write++] = heap.poll();
}

Dry run on Example 1 (k = 3):

Draining the heap then writes 8, 9, 10 into positions 4 to 6, giving [2, 3, 5, 6, 8, 9, 10].

Complexity. Time O(n log k): every element enters and leaves a heap of at most k + 1 elements once. Space O(k) for the heap.

Edge cases and pitfalls

• Window size is k + 1, not k. The occupant of position w may be at index w + k, so the heap must have seen k + 1 candidates before the first poll. Polling at size k fails on [2, 1] with k = 1, where it would output 2 first.
• k = 0. The array is already sorted. The heap version polls each element immediately; insertion sort makes no shifts.
• k ≥ n. Nothing is polled during the scan, and the final drain sorts everything: correct, just an ordinary heap sort.
• Empty or single-element arrays need no work; all three methods handle them without special cases.
• The guarantee matters. If an element is actually further than k from home, the heap method silently produces a wrong order (the tests only generate genuine k-sorted inputs). Insertion sort and Arrays.sort stay correct, merely slower.
• Duplicates are harmless: equal values are interchangeable in the output.

Related problems

• Insertion Sorting explains the shifting loop used in Approach 2.
• Merge K sorted Lists uses a heap of size k holding one candidate per list, the same “smallest pending candidate” idea.
• Merge K Sorted Arrays is the array version of that merge.
• Kth smallest element in an array [use priority queue] keeps a bounded heap to answer a selection query.
• Heaps (Theory Video) covers the heap operations and their costs.