Heap Sort in Python - Tail-recursion and state change :)

heap_sort_tail_recursion.py

# Copyright (c) 2021 Alexandru Manaila. All Rights Reserved.
# 2021/11/10

# Python is not the best for this, but the syntax makes it worth it.
# Set state_logging to false if you would rather see tests instead of state changes.
# heap_sort, fix_heap_up, and fix_heap_down are tail-recursive.
#  They reduce their parameters and re-call themselves until they can return an answer.

try_me = [1,3,2,5,4]
state_logging = True

def main():
    if state_logging:                                                 # display state change example
        print(f'\ni:  j:  [heap] - [lst]:')
        print('-----------------------')
        print('\n', heap_sort(try_me))
    else:
        print(test(sorted, heap_sort, []))                            # algorithm test cases
        print(test(sorted, heap_sort, [1]))
        print(test(sorted, heap_sort, [2, 1]))
        print(test(sorted, heap_sort, [1, 2]))
        print(test(sorted, heap_sort, [3, 2, 1]))
        print(test(sorted, heap_sort, [1, 3, 2]))
        print(test(sorted, heap_sort, [1, 2, 3]))
        print(test(sorted, heap_sort, [3, 2, 4, 1]))
        print(test(sorted, heap_sort, [5, 7, 2, 3, 4, 6, 1]))



# parameters: 
#   f, g: functions
#   t: test case
# return: formatted result string comparing f(t) and g(t) results
def test(f, g, t):
    return f'{"passed" if f(t) == g(t) else "failed"}: {t}'



# parameters:
#   lst: a list of numbers, or empty
#       'this' list will also hold the heap at the same time
#   i: optional, index for growing heap, tracks function state
#   j: optional, index for last heap element, tracks function state
# return: sorted list of numbers, or empty (eventually)
def heap_sort(lst, i=0, j=0):
    n = len(lst)
    if state_logging:
        print(f'{i},  {j},  {heap_string(lst, i, j)}')
    if i == n:                                                      # if the heap is fully built
        if j == 0:                                                  # if the heap is consumed
            return lst                                              # return the answer
        else:
            return heap_sort(fix_heap_down(pop_heap(lst, j), j), i, j-1)    # continue unraveling the heap
    else: 
        return heap_sort(fix_heap_up(lst, i), i+1, i)               # continue building the max-heap



# parameters:
#   lst: list of numbers
#   i, j: indexes of elements to be swapped
# return: new list with elements i and j swapped
def swap(lst, i, j):
    return [*lst[:i], lst[j], *lst[i+1:j], lst[i], *lst[j+1:]]      # create a list with 2 swapped elements



# parameters:
#   lst: list of numbers, should contain a heap too
#   j: index of last heap element
# return: new list with first and last heap elements swapped
def pop_heap(lst, j):
    return [lst[j], *lst[1:j], lst[0], *lst[j+1:]]          # remove top element from heap by swapping it
                                                            #  with the last heap element


# parameters:
#   lst: list of numbers, should contain a heap too
#   i: index of element being fixed upwards
# return: new list with a balanced heap (eventually)
def fix_heap_up(lst, i):
    if i == 0:                                              # if there are no more elements to test
        return lst                                          # return the answer
    else:
        p = int((i+1) / 2 - 1)                              # find parent: heap_index = lst_index + 1
        if lst[p] < lst[i]:                                 # if the parent is smaller than the child
            return fix_heap_up(swap(lst, p, i), p)          # swap the elements and continue testing upwards
        else:
            return lst                                      # return the answer



# parameters:
#   lst: list of numbers, should contain a heap too
#   n: length of heap in lst
#   i: optional, index of element being fixed downwards
# return: new list with a balanced heap (eventually)
def fix_heap_down(lst, n, i=0):
    if n <= 1 or i >= n:                                    # if the heap is too small or index is out of bounds
        return lst                                          # return the answer
    c = int((i+1) * 2 - 1)                                  # find first child, then second is c+1
    if c < n and lst[i] < lst[c]:                           # if first child in bounds and larger than parent
        if c+1 < n and lst[c] < lst[c+1]:                       # if second child in bounds and larger than first
            return fix_heap_down(swap(lst, i, c+1), n, c+1)     # swap parent with second child and continue down
        else:
            return fix_heap_down(swap(lst, i, c), n, c)         # swap parent with first child and continue down
    elif c+1 < n and lst[i] < lst[c+1]:                         # if second child in bounds and larger than parent
        return fix_heap_down(swap(lst, i, c+1), n, c+1)         # swap parent with second child and continue down
    else:
        return lst                                          # return answer since no op necessary



def heap_string(lst, i, j):
    h = i if i < len(lst) or i+j == 9 else j+1
    return f'{lst[:h]} - {lst[h:]}'



if __name__ == '__main__':
    main()