Created
September 7, 2021 10:38
-
-
Save jonaslsaa/41efd282cc50aa8ca2397200e62cc985 to your computer and use it in GitHub Desktop.
Python implementation of AVL self-balancing BST
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
# AVL Tree by vox and kritjo | |
class AVLTree: | |
root = None | |
def left_rotate(self, z): | |
y = z.right | |
T = y.left | |
y.left = z | |
z.right = T | |
z.height = 1 + max(self.height(z.left), self.height(z.right)) | |
y.height = 1 + max(self.height(y.left), self.height(y.right)) | |
return y | |
def height(self, node): | |
if node is None: | |
return 0 | |
return max(self.height(node.left), self.height(node.right)) | |
def right_rotate(self, z): | |
y = z.left | |
T = y.right | |
y.right = z | |
z.left = T | |
z.height = 1 + max(self.height(z.left), self.height(z.right)) | |
y.height = 1 + max(self.height(y.left), self.height(y.right)) | |
return y | |
def balance_factor(self, v): | |
if v is None: | |
return 0 | |
return self.height(v.left) - self.height(v.right) | |
def balance(self, v): | |
bal_factor = self.balance_factor(v) | |
if bal_factor < -1: | |
if self.balance_factor(v.right) > 0: | |
v.right = self.right_rotate(v.right) | |
return self.left_rotate(v) | |
if bal_factor > 1: | |
if self.balance_factor(v.left) < 0: | |
v.left = self.left_rotate(v.left) | |
return self.right_rotate(v) | |
return v | |
def insert(self, key, val, v=root): | |
if v is None: | |
v = self.Node(key, val) | |
elif key < v.key: | |
v.left = self.insert(key, val, v.left) | |
elif key > v.key: | |
v.right = self.insert(key, val, v.right) | |
v.height = 1 + max(self.height(v.left), self.height(v.right)) | |
return self.balance(v) | |
def remove(self, key, v=root): | |
if v is None: | |
return None | |
if key < v.key: | |
v.left = self.remove(key, v.left) | |
elif key > v.key: | |
v.right = self.remove(key, v.right) | |
elif v.left is None: | |
v = v.right | |
elif v.right is None: | |
v = v.left | |
else: | |
u = self.find_min(v.right) | |
v.key = u.key | |
v.value = u.value | |
v.right = self.remove(u.key, v.right) | |
v.height = 1 + max(self.height(v.left), self.height(v.right)) | |
return self.balance(v) | |
def find_min(self, v): | |
if v.left is None: | |
return v | |
return self.find_min(v.left) | |
class Node: | |
left = None | |
right = None | |
height = None | |
def __init__(self, key, val): | |
self.key = key | |
self.value = val |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment