AVL tree implementation in Python

Question

Is my implementation of AVL tree correct? It is my second Python program and I am not sure if do code properly.

'''
Created on 30 sty 2016

AVL tree - my second python program :)

@author: stanek
'''

class Node:

    def __init__(self, data, left = None, right = None, height = 0):
        self.data = data
        self.left = left
        self.right = right
        self.height = height

    def add_node(self, data):
        if self.data > data:
            '''Adding to left subtree'''
            if self.left is None:
                self.left = Node(data)
            else:
                self.left.add_node(data)
                if ( self.get_left_height() - self.get_right_height() == 2):
                    ''' Then we need to balance a subtree'''
                    print("Rebalancing after inserting", data)
                    if (data < self.left.data):
                        self.rotate_left()
                    else:
                        self.double_rotate_left()

        else:
            '''Adding to right subtree'''
            if self.right is None:
                self.right = Node(data)
            else:
                self.right.add_node(data)
                if ( self.get_right_height() - self.get_left_height() == 2):
                    ''' Then we need to balance a subtree'''
                    print("Rebalancing after inserting", data)
                    if (data < self.right.data):
                        self.rotate_right()
                    else:
                        self.double_rotate_right()

    def print_nodes(self):
        if self.left is not None:
            self.left.print_nodes()
        print(self);
        if self.right is not None:
            self.right.print_nodes()


    '''
    AVL methods
    '''
    def get_right_height(self):
        if self.right is None:
            return -1
        else:
            return self.right.height

    def get_left_height(self):
        if self.left is None:
            return -1
        else:
            return self.left.height

    def set_height(self):
        self.get_height = max( self.left.get_height, self.right.get_height ) + 1

    def rotate_right(self):
        print("rotating right",self.data);
        temp = self
        self = self.right
        self.right = temp
        self.right.set_height
        self.set_height

    def rotate_left(self):
        print("rotating left",self.data);
        temp = self
        self = self.left
        self.left = temp
        self.left.set_height()
        self.set_height()

    def double_rotate_right(self):
        print("double rotating right",self.data);
        temp = self.left
        self.left = temp.right
        temp.right = self
        self = temp
        self.right.set_height
        self.set_height

    def double_rotate_left(self):
        print("double rotating left",self.data);
        temp = self.right
        self.right = temp.left
        temp.left = self
        self = temp
        self.left.set_height
        self.set_height

    def __str__(self):
        return str(self.data)

class binary_tree:

    def __init__(self):
        self.root = None

    def getRoot(self):
        return self.root

    def add(self, data):
        if self.root is None:
            self.root = Node(data)
        else:
            self.root.add_node(data)

    def print_all(self):
        if self.root is None:
            print("Empty tree")
        else:
            self.root.print_nodes();


if __name__ == '__main__':
    b = binary_tree();
    b.add(1)
    b.add(2)
    b.add(4)
    b.add(6)
    b.add(3)
    b.add(5)
    b.print_all()

Answer 1

Bug

No, it is not correct:

self.right.set_height

set_height is not run due to the lack of parenthesis, hence the height is not correctly set.

Swapping

In many places such as for example:

    temp = self
    self = self.right
    self.right = temp
    self.right.set_height
    self.set_height

You are swapping values in an obfuscated manner, just use the built-in a, b = b, a syntax:

>>> a, b = 1, 2
>>> a, b = b, a
>>> a, b
(2, 1)

Naming

Universally, names of classes are written in PascalCase, and you should really stick to it, even if the compiler does not enforce it:

binary_tree -> BinaryTree

getters

Getters are highly frowned upon in Python. Values are generally accessed directly.

If you really do not want the user to access values directly you may make them "private" by prefixing __ and letting a @property return them.

Decide for your nil values

None is the standard nil (nothing) value in Python but in get_right_height you return -1 if the value isNone. If so just make the default -1 and remove the getter as suggested above.

Repetition Repetition

You repeat b.add 6 times, use a loop to avoid such waste.

---

There must be a way to avoid repetition in the below code-blocks, the problem is that you have duplicate functions for left and right, maybe a single rotate function could help you here.

        '''Adding to left subtree'''
        if self.left is None:
            self.left = Node(data)
        else:
            self.left.add_node(data)
            if ( self.get_left_height() - self.get_right_height() == 2):
                ''' Then we need to balance a subtree'''
                print("Rebalancing after inserting", data)
                if (data < self.left.data):
                    self.rotate_left()
                else:
                    self.double_rotate_left()

---

        '''Adding to right subtree'''
        if self.right is None:
            self.right = Node(data)
        else:
            self.right.add_node(data)
            if ( self.get_right_height() - self.get_left_height() == 2):
                ''' Then we need to balance a subtree'''
                print("Rebalancing after inserting", data)
                if (data < self.right.data):
                    self.rotate_right()
                else:
                    self.double_rotate_right()