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'