Friday, January 4, 2013

Heap (Max) in Python

Lately I've been practicing my data structures and algorithms and I came up with this interesting implementation for a max heap:

class Heap(object):
    """ A max heap, parents have two or less equal or smaller children """
    def __init__(self):
        """ Set the count to 0 and it has no root node """
        self.count = 0
        self.root = None

    def heapify(self, iterable):
        """ loops elements in the iterable and adds them to the heap """
        for elem in iterable:
            self.add_item(elem)

    def add_item(self, ordinal):
        """ Adds an ordinal item using priority """
        self.count += 1
        if self.count == 1:
            self.root = Node(ordinal)
            return

        i = self.count
        current_node = self.root
        new_node = Node(ordinal)

        while i != 1:
            if current_node.data < new_node.data:
                current_node.data, new_node.data = new_node.data, current_node.data
            if i%2 == 0:
                if current_node.left:
                    current_node = current_node.left
                else:
                    current_node.left = new_node
            else:
                if current_node.right:
                    current_node = current_node.right
                else:
                    current_node.right = new_node
            i /= 2

class Node(object):
    """ Used to traverse the heap """
    def __init__(self, ordinal):
        """ Require data, make the pointers empty (to None) """
        self.data = ordinal
        self.left = None
        self.right = None 


This design of a heap is pretty interesting because it uses the binary representation of Heap.count (the number of elements in the heap) to determine where the next node should go. The algorithm for this is pretty straight forward:

Increment the `count` by 1         // We're adding a new item 
`i` is `count`
`current node` is `root node`

Check 1: if `new node's value` > `current node's value`,
             then swap(`new node's value`, `current node's value`)

Check 2: if binary(`i`) ends in 0, // AKA is even, (i%2 == 0),
             then `current node` is `current node's left`
         otherwise `current node` is `current node's right`

Check 3: if `current node` is null
             then `current node` is `new node`
         otherwise `i` /= 2 and goto Check 1

The end result is a balanced heap! I also added a __str__ method to the Node class for testing:

    def __str__(self):
        return "[%s <- %s -> %s]" % \
                 (str(self.left) if self.left != None \
                                 else 'END', str(self.data), \
                  str(self.right) if self.right != None \
                                 else 'END')

No comments:

Post a Comment