Sort a K-Sorted Array Python Solution

Last Updated : 17 Mar 2025

In this problem, we will be given an array of integers. The array will be k sorted. A k-sorted array is one in which every element of the array is at most k steps away from the final sorted array, which is the target sorted array.

Let us see some examples to understand the problem

Examples:

Input: array = [5, 3, 2, 10, 9], K = 3

Output: array = {2, 3, 5, 9, 10}

Input: array = [10, 9, 7, 4, 60, 50], k = 4

Output: array = [4, 7, 9, 10, 50, 60]

Approach - 1

We will use the insertion sort algorithm to solve the problem in this approach. In this approach, the insertion sorting will place each element of the array in the correct place, hence sorting the array. This is an efficient method to solve this problem because we can use the index of every element, which the maximum k number of indices can change.

These are the steps that we will follow to solve this problem.

We will start by iterating the array from the first index to the last index.
Then, in each iteration, we will compare the element at the current index to the previous element.
If the current element is smaller than the previous element, then we will compare this element to the element before the previous element. We will continue this until we find a smaller predecessor than the current element. Once we find the index of the predecessor, we will insert the current element after this predecessor element.
We will continue this with each element of the array. At the end, we will get the final sorted array.

Below is the Python code showing how to use insertion to solve this problem.

Code

# Python program to sort the k-sorted array using the insertion sort

# Defining the function to sort the problem
def insertion(arr, n):

  # Initalizing the indices of the elements of the array
  i, c, j = 0, 0, 0

  # Starting a loop to iterate over all the elements of the array
  for i in range(n):

    # Storing the current element of the array
    c = arr[i]

    # Storing the index of the previous element
    j = i-1

    # Starting another loop to move all the elements that are before the index i to one position until we find the position at which the previous element is smaller than the current element
    # Since the array is k sorted, this loop will run at most k times.
    while j >= 0 and arr[j] > c:
      arr[j + 1] = arr[j]
      j = j - 1
    arr[j + 1] = c
  print(arr)
# Calling the function and passing the array to it
arr = [10, 9, 7, 4, 60, 50]
insertion(arr, len(arr))

Output:

[4, 7, 9, 10, 50, 60]

Time Complexity: The time complexity of this approach is non-linear since we are using nested loops. The inner loop will traverse for a maximum of k times. The outer loop will run for every element of the array. Hence, it will run for N times. Therefore, the final time complexity of this approach will be O(N * k).

Auxiliary Space: We have not used any extra space to solve this problem; therefore, the space complexity of this approach is O(1).

Approach - 2

In the second approach, we will consider the property of the k-sorted array. Earlier, we shifted every element to the right whenever the current element was not in the correct order. This is taking extra time because an element shifted x times to the left can return to the same position again if x elements are added on the left of the element. Therefore, in this approach, we will only shift the elements if the current element is more than k positions away from the correct position. This approach will optimize the time complexity of the program. This approach works well because the given array will be almost sorted; hence, we don't need to put every element in its correct order. Already, k operations of sorting are performed.

Below is the implementation of this approach.

Code

# Python program to sort the k-sorted array

# Defining the function that will take the array, the size of the array and the integer representing the value of k
def sortk(arr, n, k):

  # Starting the loop to iterate over every element of the array
  for i in range(1, n):
    curr = arr[i]
    j = i - 1

    # Now, we will check if the previous element of the current element is greater than the current element and then shift the elements to the right only if the condition on the element being more than k steps away from the correct position is true.
    while j>=max(0, i - k) and arr[j] > curr:
      arr[j + 1] = arr[j]
      j -= 1
    arr[j + 1] = curr

  # Printing the sorted array
  for i in range(n):
    print(arr[i], end = ' ')

  print()

# Calling the function and passing the array to it
arr = [10, 9, 7, 4, 60, 50]
sort(arr, len(arr), 3)

Output:

4 7 9 10 50 60

Time complexity: The time complexity of this approach is better than the previous approach. However, the worst-case time complexities of both approaches are the same. The time complexities are the same because there can be a case where each element is k distance away from the correct position. In that case, the inner loop will run for at most k times. Hence, the worst time complexities of both approaches are the same, i.e., O(N * k).

Auxiliary Space: We have not used any extra space to solve this problem; therefore, the space complexity of this approach is O(1).

Approach - 3

In this approach, we will use a heap to solve this problem.

We will follow these steps to solve the problem.

We will start by initializing a min heap, which will be of size k + 1. We are creating the heap of size k + 1 because the element of the given array can be at a maximum k distance away from the index of the correct position.
Then, we will use another loop to remove the minimum element from the heap and then put this element in the final sorted array.

Below is the implementation of this approach in Python.

Code

# Python program to sort the k-sorted array using the heap data structure

# Importing the required modules
from heapq import heappop, heappush, heapify

# Defining the function to print the array
def print_arr(arr):
  for i in arr:
    print(i, end = ' ')
  
# Defining the function that will take a k-sorted array, the size of the k-sorted array and the value of k
def sort_k(arr, n, k):
  
  # Initializing a heap containing the first k + 1 elements of the given array
  heap = arr[ : k + 1]

  # Using the heapify function to convert the Python list into a heap data structure
  heapify(heap)

  # We will run a loop over the array's remaining element, i.e., starting the k + 2nd element of the array. Our task is to place this element in the current position
  t = 0

  for i in range(k + 1, n):
    arr[t] = heappop(heap)
    heappush(heap, arr[i])
    t += 1

  while heap:
    arr[t] = heappop(heap)
    t += 1

# Calling the function and passing the array to it
arr = [10, 9, 7, 4, 60, 50]
sort_k(arr, len(arr), 3)
print_arr(arr)

Output: