Heap Sort in Python

By Gaurav Jain on February 24, 2018

As per wikipedia, Steps for Heapsort algorithm are -

  1. Convert Array/List into a max-heap structure in O(n) operations. Thats called buildMaxHeap() or heapify().
  2. After first step we’ll have max value at the root (0th index). Swap this value with last element and freez the last element. Now reduce the considered range of the list by 1.
  3. Call the siftDown() function on the list to sift the new first element to its appropriate index in the heap.
  4. Go to step2 and repeat untill considered range of the list is 1.

Visualization can be seen here

def swap(a, b):
    a, b = b, a


def heap_sort(alist):
    length = len(alist)
    # Build the heap so that largest value is at the root
    heapify(alist)

    # The following loop maintains the invariants that alist[0:end] is a heap and every element
    # beyond end is greater than everything before it, so alist[end:count] is in sorted order
    end = length - 1
    while end > 0:
        # alist[0] is the root and largest value. The swap moves it in front of the sorted elements.
        swap(alist[end], alist[0])

        # the heap size is reduced by one
        end = end - 1

        # the swap ruined the heap property, so restore it
        sift_down(alist, 0, end)

Above heap_sort() function is structured based on aforementioned rules. Now, Let’s look at the implementation of heapify and sift_down functions.


# helper functions
parent = lambda i: floor((i-1) / 2)
left_child = lambda i: 2 * i + 1
right_child = lambda i: 2 * i + 2

# Put elements of 'alist' in heap order, in-place
def heapify(alist):
    # start is assigned the index in 'a' of the last parent node
    # the last element in a 0-based array is at index length-1; find the parent of that element
    length = len(alist)
    start = parent(length - 1)

    while start >= 0:
        # sift down the node at index 'start' to the proper place such that all nodes below
        # the start index are in heap order
        sift_down(alist, start, length - 1)
        # go to the next parent node
        start = start - 1
    # after sifting down the root all nodes/elements are in heap order

# Repair the heap whose root element is at index 'start', assuming the heaps rooted at its children are valid
def sift_down(alist, start, end):
    root = start

    while left_child(root) <= end:    # While the root has at least one child
        child = left_child(root)   # Left child of root
        swap = root                # Keeps track of child to swap with

        if alist[swap] < alist[child]
            swap = child
        # If there is a right child and that child is greater
        if child + 1 <= end and alist[swap] < alist[child + 1]
            swap = child + 1
        if swap = root
            # The root holds the largest element. Since we assume the heaps rooted at the
            # children are valid, this means that we are done.
            return
        else
            swap(alist[root], alist[swap])
            root = swap            # repeat to continue sifting down the child now

This code can be found on github [Python-DSA] repository.

If you find anything wrong or has a better solution, please feel free to create issues on this repo.