""" :Author: James Sherratt :Date: 20/10/2019 :License: MIT :name: heapsort.py Heap sorts a list-like object. Note: this has been written with code-clarity in mind first, efficiency second. """ from random import randint def get_left(i): """ Get the left element index of a heap node for an array. :param i: The parent index. :return: the left element. """ return 2 * i + 1 def get_right(i): """ Get the right element index of a heap node for an array. :param i: The parent index. :return: the right element. """ return 2 * i + 2 def repair_heap(vals_list, root, arr_top): """ Sifts the root element of a heap to the correct position, to correct a max heap. This assumes the children of the root/ node are max heaps. :param vals_list: list of values, which represents a heap structure. :param root: the index of the node we're working from/ using as a root. :param arr_top: the largest value of the list we're interested in. :return: Reference to the passed list, with the root node in the correct position. """ # This is the value to swap. We want to swap the root value down, so we swap the root first. swap = root # Get left and right nodes of root. left = get_left(root) right = get_right(root) while left < arr_top: # Check if value to swap is less than the left child. if vals_list[swap] < vals_list[left]: swap = left # Check if value to swap is less than the right child (if exists). # Note: these 2 if's could be combined using "and", but then we're relying on lazy evaluation. if right < arr_top: if vals_list[swap] < vals_list[right]: swap = right # Check if the swap is still the root. If so, there's no more children to swap and we're done. if swap == root: return vals_list # Else, swap. else: vals_list[root], vals_list[swap] = vals_list[swap], vals_list[root] # New root, left and right node for the next iteration. root = swap left = get_left(root) right = get_right(root) return vals_list def max_heap(vals_list): """ Convert a list of values into a max heap tree. :param vals_list: list of numbers. :return: the same list as a max heap tree. """ # Create a max heap by repairing the heap, starting from the nodes one above the leaf nodes. len_list = len(vals_list) for root in range(len_list//2, -1, -1): repair_heap(vals_list, root, len_list) return vals_list def max_heap_to_sorted(vals_list): """ Convert a max heap list into a sorted list. :param vals_list: list containing max heap. :return: the same list of values, sorted. """ # i is the index of the last element of the slice of the array that needs sorting. for top in range(len(vals_list)-1, 0, -1): # Swap the root value (max) with the last value of the slice. vals_list[0], vals_list[top] = vals_list[top], vals_list[0] # Sift the new root to the correct position of the remainder of the max heap. # Another way of doing this is to pass a slice of the vals_list up to the value top, but python passes # slices by copy so there's a massive performance hit. repair_heap(vals_list, 0, top) return vals_list def heapsort(vals_list): """ Sort a list of values using heapsort. :param vals_list: list of sortable values. :return: the same list, sorted. """ max_heap(vals_list) return max_heap_to_sorted(vals_list) if __name__ == "__main__": list_len = 100000 vals_list = [randint(0, (2**16)) for i in range(list_len)] heap_sorted = heapsort(list(vals_list)) py_sorted = sorted(vals_list) print("Did the sort work? {}".format(heap_sorted == py_sorted))