From 6.006 Wiki
Heaps and Heapsort Implementation
# Recitation 08
# Date: 10/01/08
#
# These functions will help sort the list A[1:] using heapsort.
# Element A[0] may get modified during these operations since we
# will be storing the heap size in that slot.
###Private Helpers
def parent(i):
"""Return the parent index of i or None"""
if i/2 > 0:
return i/2
return None
def left(i):
"""Return the left child index of i"""
return 2*i
def right(i):
"""Return the right child index of i"""
return 2*i+1
def heap_size(A):
"""Returns the heap size of A"""
return A[0]
def set_heap_size(A, size):
"""Sets the heap size of A to be size"""
A[0] = size
###Public Helpers
def swap(A, i, j):
"""Swaps elements A[i] and A[j]"""
A[i], A[j] = A[j], A[i]
def is_heap_valid(A):
"""Checks if the heap invariant is met"""
for i in range(2, heap_size(A)):
if A[i] > A[parent(i)]:
return False
return True
###Heap functions
def max_heapify(A, i):
"""
Given that the children of i are heaps, create a heap rooted at i
"""
largest = i
if left(i) <= heap_size(A) and A[i] < A[left(i)]:
largest = left(i)
if right(i) <= heap_size(A) and A[largest] < A[right(i)]:
largest = right(i)
if i != largest:
swap(A, i, largest)
max_heapify(A, largest)
def extract_max(A):
"""Returns and removes the largest element from A"""
if heap_size(A) == 0:
return None
max = A[1]
A[1] = A[heap_size(A)]
set_heap_size(A, heap_size(A)-1)
max_heapify(A, 1)
return max
def build_max_heap(A):
"""Creates a max heap from list A[1:len(A)]"""
set_heap_size(A, len(A)-1)
for i in range(heap_size(A)/2, 0, -1):
max_heapify(A, i)
def heapsort(A):
"""Sorts a 1-indexed list A into ascending order"""
build_max_heap(A)
for i in range(heap_size(A), 1, -1):
A[i] = extract_max(A)
def update_value(A, i, value):
"""Sets A[i] to be value and restores heap invariant"""
if A[i] < value:
A[i] = value
float_up(A, i)
else:
A[i] = value
max_heapify(A, i)
def float_up(A, i):
"""Swaps A[i] with its parent until A[i] < A[parent(i)]"""
if parent(i) and A[parent(i)] < A[i]:
swap(A, i, parent(i))
float_up(A, parent(i))
# Create some sample lists
A = [20, 10, 17, 19, 5, 1, 9, 14, 18, 6, 20, 3, 15, 2, 11, 8, 13, 4, 12, 7, 16]
A1 = [None, 10, 17, 19, 5, 1, 9, 14, 18, 6, 20, 3, 15, 2, 11, 8, 13, 4, 12, 7, 16]
B = [6, 'cat', 'bat', 'ant', 'dog', 'emu', 'frog']
B1 = [None, 'cat', 'bat', 'ant', 'dog', 'emu', 'frog']
