Problem

You are given a list of N numbers, in which each number is located at most k places away from its sorted position.

For example, if k = 3, a given element at index 4 might end up at any index between [1, 7]

1
2
3
indices: 0 1 2 3 4 5 6 7 8 9 10 11
           ^     ^     ^
           <---------->	

Come up with an algorithm that sorts this list in O(N log k) time.

Examples

Example 1

1
2
3
Input: nums = [6, 5, 3, 2, 8, 10, 9], k = 3
Output: [2, 3, 5, 6, 8, 9, 10]
Explanation: Each element is at most 3 positions away from its sorted position.

Example 2

1
2
3
Input: nums = [3, 2, 1, 5, 4, 7, 6, 5], k = 1
Output: [1, 2, 3, 4, 5, 5, 6, 7]
Explanation: Each element is at most 1 position away from its sorted position.

Example 3

1
2
3
Input: nums = [1, 4, 5, 2, 3, 7, 8, 6, 10, 9], k = 2
Output: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
Explanation: Each element is at most 2 positions away from its sorted position.

Solution

Solution

Method 1 - Brute Force Sorting

The simplest approach is to ignore the nearly sorted property and apply any standard sorting algorithm. However, this doesn’t leverage the fact that elements are already close to their correct positions.

Complexity

  • ⏰ Time complexity: O(N log N), standard sorting without using k-sorted property
  • 🧺 Space complexity: O(1) for in-place sorting algorithms

Method 2 - Insertion Sort

Intuition

Insertion sort performs well on nearly sorted arrays. Since each element is at most k positions away from its correct place, we can leverage this property to achieve better performance than general sorting.

Approach

  1. Use insertion sort which naturally handles nearly sorted arrays efficiently
  2. For each element, we only need to shift it at most k positions to the left
  3. Since elements are close to their correct positions, insertion sort becomes much faster
  4. Time complexity improves from O(N²) to O(N×k) due to limited displacement

Code

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
class Solution {
public:
    vector<int> sortArray(vector<int>& nums, int k) {
        for (int i = 1; i < nums.size(); i++) {
            int key = nums[i];
            int j = i - 1;
            while (j >= 0 && nums[j] > key) {
                nums[j + 1] = nums[j];
                j--;
            }
            nums[j + 1] = key;
        }
        return nums;
    }
};
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
func sortArray(nums []int, k int) []int {
    for i := 1; i < len(nums); i++ {
        key := nums[i]
        j := i - 1
        for j >= 0 && nums[j] > key {
            nums[j+1] = nums[j]
            j--
        }
        nums[j+1] = key
    }
    return nums
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
class Solution {
    public int[] sortArray(int[] nums, int k) {
        for (int i = 1; i < nums.length; i++) {
            int key = nums[i];
            int j = i - 1;
            while (j >= 0 && nums[j] > key) {
                nums[j + 1] = nums[j];
                j--;
            }
            nums[j + 1] = key;
        }
        return nums;
    }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
class Solution:
    def sortArray(self, nums: List[int], k: int) -> List[int]:
        for i in range(1, len(nums)):
            key = nums[i]
            j = i - 1
            while j >= 0 and nums[j] > key:
                nums[j + 1] = nums[j]
                j -= 1
            nums[j + 1] = key
        return nums

Complexity

  • ⏰ Time complexity: O(N×k), each element moves at most k positions, better than O(N²)
  • 🧺 Space complexity: O(1), in-place sorting

Method 3 - Min Heap Approach

Intuition

Since each element is at most k positions away from its correct place, we can use a min heap of size k+1. The smallest element in the first k+1 elements must be the first element in the sorted array. We can continuously extract the minimum and add new elements to maintain this property.

Approach

  1. Create a min heap with the first k+1 elements
  2. Extract the minimum element (this is the next element in sorted order)
  3. Add the next element from the array to the heap
  4. Repeat until all elements are processed
  5. Extract remaining elements from the heap

Code

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
class Solution {
public:
    vector<int> sortArray(vector<int>& nums, int k) {
        priority_queue<int, vector<int>, greater<int>> minHeap;
        int n = nums.size();
        vector<int> result;
        
        // Add first k+1 elements to heap
        for (int i = 0; i <= k && i < n; i++) {
            minHeap.push(nums[i]);
        }
        
        // Process remaining elements
        for (int i = k + 1; i < n; i++) {
            result.push_back(minHeap.top());
            minHeap.pop();
            minHeap.push(nums[i]);
        }
        
        // Extract remaining elements
        while (!minHeap.empty()) {
            result.push_back(minHeap.top());
            minHeap.pop();
        }
        
        return result;
    }
};
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
import "container/heap"

type IntHeap []int

func (h IntHeap) Len() int           { return len(h) }
func (h IntHeap) Less(i, j int) bool { return h[i] < h[j] }
func (h IntHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
func (h *IntHeap) Push(x interface{}) { *h = append(*h, x.(int)) }
func (h *IntHeap) Pop() interface{} {
    old := *h
    n := len(old)
    x := old[n-1]
    *h = old[0 : n-1]
    return x
}

func sortArray(nums []int, k int) []int {
    minHeap := &IntHeap{}
    heap.Init(minHeap)
    n := len(nums)
    result := make([]int, 0, n)
    
    // Add first k+1 elements to heap
    for i := 0; i <= k && i < n; i++ {
        heap.Push(minHeap, nums[i])
    }
    
    // Process remaining elements
    for i := k + 1; i < n; i++ {
        result = append(result, heap.Pop(minHeap).(int))
        heap.Push(minHeap, nums[i])
    }
    
    // Extract remaining elements
    for minHeap.Len() > 0 {
        result = append(result, heap.Pop(minHeap).(int))
    }
    
    return result
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
class Solution {
    public int[] sortArray(int[] nums, int k) {
        PriorityQueue<Integer> minHeap = new PriorityQueue<>();
        int n = nums.length;
        int[] result = new int[n];
        int resultIndex = 0;
        
        // Add first k+1 elements to heap
        for (int i = 0; i <= k && i < n; i++) {
            minHeap.add(nums[i]);
        }
        
        // Process remaining elements
        for (int i = k + 1; i < n; i++) {
            result[resultIndex++] = minHeap.poll();
            minHeap.add(nums[i]);
        }
        
        // Extract remaining elements
        while (!minHeap.isEmpty()) {
            result[resultIndex++] = minHeap.poll();
        }
        
        return result;
    }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
class Solution:
    def sortArray(self, nums: List[int], k: int) -> List[int]:
        import heapq
        
        min_heap = []
        n = len(nums)
        result = []
        
        # Add first k+1 elements to heap
        for i in range(min(k + 1, n)):
            heapq.heappush(min_heap, nums[i])
        
        # Process remaining elements
        for i in range(k + 1, n):
            result.append(heapq.heappop(min_heap))
            heapq.heappush(min_heap, nums[i])
        
        # Extract remaining elements
        while min_heap:
            result.append(heapq.heappop(min_heap))
        
        return result

Complexity

  • ⏰ Time complexity: O(N log k), each of N elements involves heap operations of O(log k)
  • 🧺 Space complexity: O(k), heap stores at most k+1 elements