window_heapsort.py

from heapq import heappush, heappop


# This sort runs in O(n log k + k log k),
# which simplifies to O(n log k) whenever k = o(n)
# (k is sublinear in n).
def window_heapsort(a, k):
    if k > len(a):
        return sorted(a)

    window_width = k + 1
    window_heap = []
    result_array = []
    index = 0

    # Initialize the window heap:
    # Runs in O(k log k)
    for i in range(window_width):
        heappush(window_heap, a[i])

    # Keep sliding the window over the array:
    # Here, the size of the heap alternates between
    # 'window_width' and 'window_width - 1', which is
    # O(k), and thus pops and pushes run in O(log k).
    # Since we march over (at most) n elements, the
    # complexity of this loop is O(n log k).
    while index + window_width < len(a):
        result_array.append(heappop(window_heap))
        heappush(window_heap, a[window_width + index])
        index += 1

    # Dump the leftovers to the end of the result array:
    # Runs in O(k log k)
    while window_heap:
        result_array.append(heappop(window_heap))

    return result_array


def sort_almost_sorted(a, k):
    if not a:
        return []

    length = len(a)
    if k >= length:
        return sorted(a)  # apply built-in "timsort", designed to be quick on almost sorted lists

    result = []
    min_heap = []  # maintain a min heap for a sliding window
    for slice_start in range(0, length, k + 1):  # sliding window
        print("new slice")
        # push the next slice to min heap
        for index in range(slice_start, min(slice_start + k + 1, length)):
            print("inner push")
            heappush(min_heap, a[index])

        print("inner pop")
        result.append(heappop(min_heap))

    # put the rest of the heap into the result
    for _ in range(len(min_heap)):
        print("end pop")
        result.append(heappop(min_heap))

    return result

arr = [0, 3, 2, 1, 4, 5, 6, 8, 7, 9]
k = 2

result_op = sort_almost_sorted(arr, k)
result_my = window_heapsort(arr, k)

print(arr)
print(result_op)
print(result_my)