A heap implementation in Python

One more python implementation of a data structure. This is a heap implementation that follows an architecture similar to the binary search tree we have built earlier.

from collections import Container

class Heap(Container):
    array = []

    # Allows "if item in tree:"
    def __contains__(self, item):
        return self._contains(item, 1)

    def _contains(self, item, index):
        if item > self.array[index-1]:
            return False
        elif item < self.array[index-1]:
            if len(self.array) > index * 2:
                return self._contains(item, index*2) or self._contains(item, index*2+1)
            elif len(self.array) > index * 2 + 1:
                return self._contains(item, index*2)
            else:
                return False
        else:
            return True

    def add(self, item):
        self.array.append(item)
        self._bubble(len(self.array)-1)

    def _bubble(self, index):
        # Note: we take 1-based indexes, but convert them to 0-based ones
        index = index-1
        if self.array[index] > self.array[index/2]:
            tmp = self.array[index]
            self.array[index] = self.array[index/2]
            self.array[index/2] = tmp
            self._bubble(index/2)

    # Useful for debugging. Prints the whole tree
    def __repr__(self):
        return self._repr(1)

    def _repr(self, index, indent=0):
        # Note: we use 1-based indexes here
        string = ""
        string = ('\t' * indent) + str(self.array[index-1]) + '\n'

        if len(self.array) >= index * 2:
            string += ('\t' * (indent+1)) + 'Left:\n'
            string += self._repr(index*2, indent+1)

        if len(self.array) >= index * 2 + 1:
            string += ('\t' * (indent+1)) + 'Right:\n'
            string += self._repr(index*2+1, indent+1)

        return string

You can reuse the same code as the last time to test it:

from random import randint

# Add random values
values = []
tree = Heap()
for i in range(40):
    val = randint(0,100)
    tree.add(val)
    values.append(val)

print(tree)  # uses __repr__
print(values)
print(25 in tree)  # uses __contains__
print(25 in values)

One comment on “A heap implementation in Python

Leave a Reply

Your email address will not be published. Required fields are marked *

To create code blocks or other preformatted text, indent by four spaces:

    This will be displayed in a monospaced font. The first four 
    spaces will be stripped off, but all other whitespace
    will be preserved.
    
    Markdown is turned off in code blocks:
     [This is not a link](http://example.com)

To create not a block, but an inline code span, use backticks:

Here is some inline `code`.

For more help see http://daringfireball.net/projects/markdown/syntax