/Interview Study Guide/Algorithms & data structures
#121

Sort a K-Sorted Array

medium
heap-priority-queuesortingarray

You are given an integer array nums that is almost sorted: every element is at most k positions away from its correct position in the fully sorted order. Return the array fully sorted in non-decreasing order.

For example, with nums = [3, 1, 2, 5, 4] and k = 2, each value is within two slots of where it belongs, and the sorted result is [1, 2, 3, 4, 5].

The point is to do better than a general sort by exploiting k: because no element travels more than k slots, the next-smallest element is always within the first k + 1 of the unplaced elements. A min-heap of size k + 1 therefore yields the sorted order in O(n log k) time.

Example

Input: nums = [3,1,2,5,4], k = 2
Output: [1,2,3,4,5]

Each value is within two positions of its sorted slot; a size-3 min-heap pops them in order.

Constraints

  • 1 <= nums.length <= 10^5
  • 0 <= k < nums.length
  • -10^9 <= nums[i] <= 10^9
  • Each element is at most k positions away from its sorted position.

Intuition

The obvious approach ignores the k guarantee entirely and just runs a general comparison sort over the whole array.

function sortKSortedArray(nums, k) {
  // A general sort works, but does O(n log n) regardless of how small k is.
  return [...nums].sort((a, b) => a - b);
}
Brute force — a full sort, ignoring the k bound: O(n log n).

A full sort is O(n log n) and never uses the promise that no element is more than k slots from home. Can we do better?

The key observation: because every element is within k positions of its sorted spot, the smallest unplaced element is always among the next `k + 1` elements. So keep a min-heap of size `k + 1`: seed it with the first k + 1 elements, then for each remaining element pop the heap minimum (the next sorted value) and push the newcomer. Drain the heap at the end. This is the bounded-heap idea from the Heaps intro, sized to the window the guarantee gives us.

Each of the n elements does one O(log k) push and one O(log k) pop — O(n log k), beating the full sort when k ≪ n.

Walking through nums = [4, 2, 1, 3, 6, 5], k = 2 (heap of size k + 1 = 3):

nums = [4, 2, 1, 3, 6, 5], k = 2 — min-heap of size k+1 = 3

40
21
in
12
33
64
55
seed heap {4, 2, 1} → root 1

Push the first k+1 = 3 elements. The global minimum (1) must be among them — root is 1.

40
21
12
in
33
64
55
pop 1, push 3 → {4, 2, 3} → output [1]

Pop the min (1) as the first sorted value, then push the next element, 3. Output: 1.

40
21
12
33
in
64
55
pop 2, push 6 → {4, 6, 3} → output [1,2]

Min is now 2 — pop and output it, push 6. Output: 1, 2.

40
21
12
33
64
in
55
pop 3, push 5 → {4, 6, 5} → output [1,2,3]

Pop 3, push the last element 5. Output: 1, 2, 3. No elements remain to scan.

40
21
12
33
64
55
drain {4, 5, 6} → output [1,2,3,4,5,6]

Drain the heap in order: 4, 5, 6. Final sorted array: [1, 2, 3, 4, 5, 6].

Optimization

Size-(k+1) min-heap

Because every element is at most k positions from its sorted spot, the smallest unplaced element is always among the first k + 1 elements still in play. Maintain a min-heap of size k + 1: push the first k + 1 elements, then for each remaining element pop the heap minimum (the next value in sorted order) and push the new element. After the scan, drain the heap.

A binary min-heap with sift-up / sift-down gives O(log k) push and pop, and we do O(n) of each — O(n log k) time. The heap holds at most k + 1 elements, so O(k) extra space.

function sortKSortedArray(nums, k) {
  // Min-heap as a binary heap over an array; heap[0] is always the minimum.
  const heap = [];
  const swap = (i, j) => { const t = heap[i]; heap[i] = heap[j]; heap[j] = t; };
  const up = (i) => {
    while (i > 0) {
      const parent = (i - 1) >> 1;
      if (heap[parent] <= heap[i]) break;
      swap(parent, i);
      i = parent;
    }
  };
  const down = (i) => {
    const n = heap.length;
    while (true) {
      let smallest = i;
      const l = 2 * i + 1;
      const r = 2 * i + 2;
      if (l < n && heap[l] < heap[smallest]) smallest = l;
      if (r < n && heap[r] < heap[smallest]) smallest = r;
      if (smallest === i) break;
      swap(i, smallest);
      i = smallest;
    }
  };
  const push = (x) => { heap.push(x); up(heap.length - 1); };
  const pop = () => {
    const top = heap[0];
    const last = heap.pop();
    if (heap.length > 0) { heap[0] = last; down(0); }
    return top;
  };

  const result = [];
  const limit = Math.min(k + 1, nums.length);
  // Seed the heap with the first k+1 elements: the global minimum must be among them.
  for (let i = 0; i < limit; i++) push(nums[i]);
  // For each remaining element, the heap min is the next sorted value; swap it out for the newcomer.
  for (let i = limit; i < nums.length; i++) {
    result.push(pop());
    push(nums[i]);
  }
  // Drain whatever is left in sorted order.
  while (heap.length > 0) result.push(pop());
  return result;
}

Complexity analysis

Time complexity: O(n log k). Here's why:

  • The heap is bounded at k + 1 elements, so every push and pop is O(log k).
  • Each of the n elements is pushed once and popped once.

So the work is n × O(log k) — overall O(n log k), beating the O(n log n) full sort when k ≪ n.

Space complexity: O(k). Here's why:

  • The heap holds at most k + 1 elements at any moment.

So the extra space is O(k) (not counting the output array).

Test cases

Beyond the example above, these are worth thinking through before you submit.

InputExpected outputDescription
nums = [1], k = 0[1]Single element, k = 0 — already in place.
nums = [9,8], k = 1[8,9]One adjacent swap away from sorted; a size-2 heap fixes it.
nums = [3,3,3], k = 1[3,3,3]All equal — order is stable and unchanged.
nums = [0,-1,-2], k = 2[-2,-1,0]Negatives, fully reversed within the k = 2 window.
nums = [2,4,1,3,5], k = 2[1,2,3,4,5]Each value within two slots of home; a size-3 heap pops them in order.

Try it yourself

Write your solution against the real judge before checking the reference.

Open in editor