Sort k-messed array

kmessed_sort.py

import heapq as hq

def sort_k_messed(a, k):
    """Sort an array where each element is at most k steps away from 
    its sorted position.
    
    Parameters:
    -----------
    a: list
        k-messed input array
    k: int
        maximum number of steps each element is away from its position
    
    Returns:
    --------
    
    a: list
        sorted array
    
    Algorithm:
    ---------
    
    Algo#1: generic sort, e.g., merge sort (O(n*log(n)) time, O(n) space) or quick sort (O(n*log(n)), O(log(n) space))
    Algo#2: if integers are within a range, bucket\radix sort -> ~O(n+k), where k is the number of buckets
    Algo#3: insertion-sort based? Take a slice, find minimum element in O(k), place it. Make n steps -> O(n*k)
    Algo#4: 
    
    1. Take a[0:k+1] buffer, form min-heap. Set next_arr_ix=k+2, sorted_ix=0;
    2. For all remaining elements of the array:
        2.1. Get min element from heap, place it at position a[sorted_ix], sorted_ix++
        2.2. Get next element from arr, place it to min heap
    3. For all remaining elements in the heap, get min element, place it at the position a[sorted_ix], sorted_ix++
    
    Complexity
    ----------
    Space: O(k)
    Time: O(n*log(k)), since place\extract element from min heap - O(log(k)), we need to do it n times 
    
    
    Example: 
    [1, 4, 5, 2, 3]
     *
    
    1. [1, 4, 5]
    min -> 1
    2. [4, 5, 2]
    min -> 2
    3. [4, 5, 3]
    min -> 3
    
    Left: [4, 5]
    min -
    
    Result: [1, 2, 3, 4, 5]
    
    
    Corner cases:
    ------------
    1. zero-sized array
    2. k > n
    3. duplicate elements
    

    """
    
    # create initial heap
    frame = list(a[0:k+1])
    hq.heapify(frame)
    sorted_ix = 0
    
    # process rest of the array
    for ix in range(k+1, len(a)):
        a[sorted_ix] = hq.heappop(frame)
        hq.heappush(frame, a[ix])
        sorted_ix += 1
    
    # process rest of the heap
    for ix in range(0, len(frame)):
        a[sorted_ix] = hq.heappop(frame)
        sorted_ix+=1
        
    return a