Recitation08 Notes

From 6.006 Wiki

Latest revision as of 15:17, 7 October 2008

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']