# AVL Trees¶

## Rules¶

1. If a subtree needs a left rotation to bring it to balance:

• First, check the balance factor of the right child.
• If the right child is left heavy, then do a right rotation on the right child
• Peform the original left rotation
2. If a subtree needs a right rotation to bring it to balance:
• First, check the balance factor of the left child.
• If the left child is right heavy, then do a left rotation on the left child
• Peform the original right rotation
In [1]:
from pythonds9.tree.bst import TreeNode, BinarySearchTree
import requests

In [2]:
def get_names(url):
req = requests.get(url)
lines = req.text.split("\n")
names = []

for line in lines:
data = line.split()
if len(data) == 4:
names.append( (data[0], data[3]) )

return names

In [3]:
def create_graph():
names = get_names("https://www2.census.gov/topics/genealogy/1990surnames/dist.female.first")
print(f"Extracted {len(names)} names")

# names[("MARY", "1"), ...)
# Sort the names list by name, not rank
names.sort(key=lambda t: t[0])

name_tree = AVLBinarySearchTree()

for name, rank in names:
name_tree[name] = int(rank)

return name_tree


## AVL TreeNode through subclassing¶

In [4]:
def foo(*args, **kwargs):
print(args, kwargs)

a = ["foo", 6, 3.14159]
k = {"bar": "chocolate", "fubar": "has bad words in it!"}
print(foo(*a, **k))

foo("foo", 6, 3.14159, bar="chocolate", fubar="has bad words in it!")

('foo', 6, 3.14159) {'bar': 'chocolate', 'fubar': 'has bad words in it!'}
None
('foo', 6, 3.14159) {'bar': 'chocolate', 'fubar': 'has bad words in it!'}

In [5]:
class AVLTreeNode(TreeNode):
"""A subclass of TreeNode"""
def __init__(self, *args, **kwargs):
# Initialize a TreeNode instance that we subclass from)
super(AVLTreeNode, self).__init__(*args, **kwargs)
self.balance_factor = 0


## AVLBinarySearchTree¶

In [6]:
class AVLBinarySearchTree(BinarySearchTree):
def put(self, key, val):
if self.root:
self._put(key, val, self.root)
else:
self.root = AVLTreeNode(key, val)

self.size += 1

def _put(self, key, val, current_node):
if key < current_node.key:
if current_node.has_left_child():
self._put(key, val, current_node.left_child)
else:
current_node.left_child = AVLTreeNode(key, val, parent=current_node)
self.update_balance(current_node.left_child)
else:
if current_node.has_right_child():
self._put(key, val, current_node.right_child)
else:
current_node.right_child = AVLTreeNode(key, val, parent=current_node)
self.update_balance(current_node.right_child)

def update_balance(self, node):
if node.balance_factor > 1 or node.balance_factor < -1:
self.rebalance(node)
return

if node.parent is not None:
if node.is_left_child():
node.parent.balance_factor += 1
elif node.is_right_child():
node.parent.balance_factor -= 1

if node.parent.balance_factor != 0:
self.update_balance(node.parent)

def rebalance(self, node):
if node.balance_factor < 0:
if node.right_child.balance_factor > 0:
self.rotate_right(node.right_child)
self.rotate_left(node)
else:
self.rotate_left(node)
elif node.balance_factor > 0:
if node.left_child.balance_factor < 0:
self.rotate_left(node.left_child)
self.rotate_right(node)
else:
self.rotate_right(node)

def rotate_left(self, rot_root):
new_root = rot_root.right_child
rot_root.right_child = new_root.left_child

if new_root.left_child is not None:
new_root.left_child.parent = rot_root
new_root.parent = rot_root.parent

if rot_root.is_root():
self.root = new_root
else:
if rot_root.is_left_child():
rot_root.parent.left_child = new_root
else:
rot_root.parent.right_child = new_root

new_root.left_child = rot_root
rot_root.parent = new_root

rot_root.balance_factor = (rot_root.balance_factor + 1
- min(new_root.balance_factor, 0))

new_root.balance_factor = (new_root.balance_factor + 1
+ max(rot_root.balance_factor, 0))

def rotate_right(self, rot_root):
new_root = rot_root.left_child
rot_root.left_child = new_root.right_child

if new_root.right_child != None:
new_root.right_child.parent = rot_root
new_root.parent = rot_root.parent

if rot_root.is_root():
self.root = new_root
else:
if rot_root.is_left_child():
rot_root.parent.left_child = new_root
else:
rot_root.parent.right_child = new_root

new_root.right_child = rot_root
rot_root.parent = new_root

rot_root.balance_factor = (rot_root.balance_factor - 1
- max(new_root.balance_factor, 0))

new_root.balance_factor = (new_root.balance_factor - 1
+ min(rot_root.balance_factor, 0))


## Visualizing AVLBinarySearchTree¶

In [7]:
nt = create_graph()

Extracted 4275 names

In [8]:
from pythonds9.tree.utils import viz_tree

In [9]:
g = viz_tree(nt)

In [10]:
g.render("avl_names")

Out[10]:
'avl_names.pdf'