Treap Insertion in Python. (Deletion is not implemented)

treap.py

from random import randint
import ipdb
RAND_MAX = 1000


class Node(object):
    def __init__(self, key, data, priority=None):
        self.key = key
        self.data = data
        
        self.priority = priority
        if priority is None:
            self.priority = randint(1, RAND_MAX)

        self.left = None
        self.right = None
        self.parent = None

    def __repr__(self):
        return "Key: {}, Priority: {}".format(self.key, self.priority)


class Treap(object):

    def __init__(self):
        self.root = None


    def insert(self, key, data, priority=None):
        node = Node(key, data, priority)
        if self.root is None:
            self.root = node
            return

        else:
            self.tree_insert(self.root, key, node)
            self.insert_heap_balance(node)

        self.visualize()


    def tree_insert(self, tree_root, key, node):

        if key <= tree_root.key:
            if tree_root.left is None:
                tree_root.left = node
                node.parent = tree_root
            else:
                self.tree_insert(tree_root.left, key, node)


        else:
            if tree_root.right is None:
                tree_root.right = node
                node.parent = tree_root
            else:
                self.tree_insert(tree_root.right, key, node)


    def visualize(self):
        
        current_nodes = [self.root]
        current_level = 1
        while len(current_nodes):
            gap = 4 - current_level
            str_nodes = (str(node) for node in current_nodes)
            ok = ""
            for n in str_nodes:
                ok += "{}{}{}".format(" "*gap, n, " "*gap)
            
            print ok
            current_level += 1
            current_nodes = self.get_children(current_nodes)


    def get_children(self, nodes):
        children = []
        for node in nodes:
            if node.left is not None:
                children.append(node.left)
            if node.right is not None:
                children.append(node.right)

        return children

    def get_height(self):
        if self.root is None:
            return 0
        
        height = 0
        current_nodes = [self.root]
        while len(current_nodes):
            height += 1
            current_nodes = self.get_children(current_nodes)

        return height

    def get(self, key):
        
        if self.root is None:
            return None

        current_node = self.root
        while current_node is not None and current_node.key != key:
            if key < current_node.key:
                current_node = current_node.left
            else:
                current_node = current_node.right

        return current_node


    def insert_heap_balance(self, node):

        iterations = 0

        #ipdb.set_trace()

        while (self.root is not node) and (node.priority >= node.parent.priority):
            iterations += 1
            if iterations > 10:
                ipdb.set_trace()

            parent = node.parent

            if node.key <= parent.key:
                if parent.parent is not None:
                    if parent.parent.left == parent:
                        parent.parent.left = node
                    elif parent.parent.right == parent:
                        parent.parent.right = node

                parent.left = node.right
                if parent.left is not None:
                    parent.left.parent = parent
                node.right = parent

            else:
                if parent.parent is not None:
                    if parent.parent.left == parent:
                        parent.parent.left = node
                    elif parent.parent.right == parent:
                        parent.parent.right = node

                parent.right = node.left
                if parent.right is not None:
                    parent.right.parent = parent

                node.left = parent

            node.parent = parent.parent
            parent.parent = node
            if node.parent is None:
                self.root = node

    def in_order(self):
        if self.root is None:
            return []

        return self.in_order_impl(self.root)

    def in_order_impl(self, tree_root):
        if tree_root.left is not None:
            self.in_order_impl(tree_root.left)

        print tree_root

        if tree_root.right is not None:
            self.in_order_impl(tree_root.right)

    def pre_order_impl(self, tree_root):
        print tree_root

        if tree_root.left is not None:
            self.in_order_impl(tree_root.left)
        
        if tree_root.right is not None:
            self.in_order_impl(tree_root.right)


    def delete(self, key):
        pass


def run():
    t = Treap()
    t.insert(4,4, 16)
    t.insert(5,5, 26)
    t.insert(6,6, 28)
    t.insert(1,1, 39)
    t.insert(2,2, 49)
    t.insert(3,3, 59)

    t.visualize()
    return t