AVL Tree in Python

Here is an AVL tree implemented in Python! Note that it is supposed to be "beautiful" and easy to read and not efficient, that is why both balance() and height() are recursive. I learnt the algorithm from C++ an Introduction to Data Structures which is an excellent book that I recommend. I hope this snippet is useful to someone:
# -*- coding: utf-8 -*- class Node: def __init__(self, data): self.data = data self.set_childs(None, None) def set_childs(self, left, right): self.left = left self.right = right def balance(self): lheight = 0 if self.left: lheight = self.left.height() rheight = 0 if self.right: rheight = self.right.height() return lheight - rheight def height(self): lheight = 0 if self.left: lheight = self.left.height() rheight = 0 if self.right: rheight = self.right.height() return 1 + max(lheight, rheight) def rotate_left(self): self.data, self.right.data = self.right.data, self.data old_left = self.left self.set_childs(self.right, self.right.right) self.left.set_childs(old_left, self.left.left) def rotate_right(self): self.data, self.left.data = self.left.data, self.data old_right = self.right self.set_childs(self.left.left, self.left) self.right.set_childs(self.right.right, old_right) def rotate_left_right(self): self.left.rotate_left() self.rotate_right() def rotate_right_left(self): self.right.rotate_right() self.rotate_left() def do_balance(self): bal = self.balance() if bal > 1: if self.left.balance() > 0: self.rotate_right() else: self.rotate_left_right() elif bal < -1: if self.right.balance() < 0: self.rotate_left() else: self.rotate_right_left() def insert(self, data): if data <= self.data: if not self.left: self.left = Node(data) else: self.left.insert(data) else: if not self.right: self.right = Node(data) else: self.right.insert(data) self.do_balance() def print_tree(self, indent = 0): print " " * indent + str(self.data) if self.left: self.left.print_tree(indent + 2) if self.right: self.right.print_tree(indent + 2) if __name__ == "__main__": tree = Node(5) tree.insert(7) tree.insert(9) tree.print_tree()
I just had to update this so that it would use COLORS!

6 kommentarer:

Anonym sa...

It indeed is useful. Thanks. :)

Anonym sa...

Yes very useful and easier to understand than most implementations I've seen. Thank you!

Jekaterina sa...

I completely dont understand how it prints the tree

Anonym sa...

Since this is the first page that comes up in Google for the search "python avl tree", I think I must leave a comment. The idea of using a recursive method for calculating height is a really, really bad one. It makes your implementation run in linear time because a call of height() at root requires going through all nodes. This completely destroys the whole point of using binary trees in the first place.

No, I didn't use this myself, but was really scared when I saw this. There might be someone out there who just graps code from the internet thinking that it must be ok. Why would anyone put faulty code in the internet ...

Anonym sa...

Really hurts to say that there is something wrong with this :(
The very first input i supplied didn't work.

The input was
14 22 27 30 33 34 35 36 40 42 55 58 78 79 80 89 95
in that order
that is,

tree = Node(14)


the result i am getting is not a balanced tree :(

Anonym sa...

I know this blog post is old, but could you please add a "delete" as well? That would be much appreciated :)