'''
PROJECT 5 - AVL Trees
Name: Lingyun Dai
PID: A54577759
'''
import random as r # To use for testing
class Node:
# DO NOT MODIFY THIS CLASS #
__slots__ = 'value', 'parent', 'left', 'right', 'height'
def __init__(self, value, parent=None, left=None, right=None):
"""
Initialization of a node
:param value: value stored at the node
:param parent: the parent node
:param left: the left child node
:param right: the right child node
"""
self.value = value
self.parent = parent
self.left = left
self.right = right
self.height = 0
def __eq__(self, other):
"""
Determine if the two nodes are equal
:param other: the node being compared to
:return: true if the nodes are equal, false otherwise
"""
if type(self) is not type(other):
return False
return self.value == other.value
def __str__(self):
"""String representation of a node by its value"""
return str(self.value)
def __repr__(self):
"""String representation of a node by its value"""
return str(self.value)
class AVLTree:
def __init__(self):
# DO NOT MODIFY THIS FUNCTION #
"""
Initializes an empty Binary Search Tree
"""
self.root = None # The root Node of the tree
self.size = 0 # The number of Nodes in the tree
def __eq__(self, other):
# DO NOT MODIFY THIS FUNCTION #
"""
Describe equality comparison for BSTs ('==')
:param other: BST being compared to
:return: True if equal, False if not equal
"""
if self.size != other.size:
return False
if self.root != other.root:
return False
if self.root is None or other.root is None:
return True # Both must be None
if self.root.left is not None and other.root.left is not None:
r1 = self._compare(self.root.left, other.root.left)
else:
r1 = (self.root.left == other.root.left)
if self.root.right is not None and other.root.right is not None:
r2 = self._compare(self.root.right, other.root.right)
else:
r2 = (self.root.right == other.root.right)
result = r1 and r2
return result
def _compare(self, t1, t2):
# DO NOT MODIFY THIS FUNCTION #
"""
Recursively compares two trees, used in __eq__.
:param t1: root node of first tree
:param t2: root node of second tree
:return: True if equal, False if not
"""
if t1 is None or t2 is None:
return t1 == t2
if t1 != t2:
return False
result = self._compare(t1.left, t2.left) and self._compare(t1.right, t2.right)
return result
def visual(self):
"""
Returns a visual representation of the AVL Tree in terms of levels
:return: None
"""
root = self.root
if not root:
print("Empty tree.")
return
bfs_queue = []
track = {}
bfs_queue.append((root, 0, root.parent))
h = self.height(self.root)
for i in range(h+1):
track[i] = []
while bfs_queue:
node = bfs_queue.pop(0)
track[node[1]].append(node)
if node[0].left:
bfs_queue.append((node[0].left, node[1] + 1, node[0]))
if node[0].right:
bfs_queue.append((node[0].right, node[1] + 1, node[0]))
for i in range(h+1):
print(f"Level {i}: ", end='')
for node in track[i]:
print(tuple([node[0], node[2]]), end=' ')
print()
### Implement/Modify the functions below ###
def insert(self, node, value):
"""
Takes in a value to be added in the form of a node to the tree
Takes in the root of the (sub)tree the node will be added into
Do nothing if the value is already in the tree
Must be recursive
Balances the tree if it needs it
O(log(n)) time complexity, O(1)* space complexity
:param node: root of the tree
:param value: the value to be added
:return: new root of the tree
"""
# check if node is root or not
if self.root is None and node == self.root:
self.root = Node(value)
self.size += 1
elif node is not None and node.value != value:
if node.value < value:
# insert to the right
if node.right:
self.size -= 1
self.insert(node.right, value)
else:
node.right = Node(value)
node.right.parent = node
elif node.value > value:
if node.left:
self.size -= 1
self.insert(node.left, value)
else:
node.left = Node(value)
node.left.parent = node
node.height = max(self.height(node.left),
self.height(node.right)) + 1
# rebalance
self.rebalance(node)
self.size += 1
def remove(self, node, value):
"""
Takes in a value to remove from the tree
Takes in the root of the (sub)tree the node will be removed from
Do nothing if the value is not found
When removing a value with two children, replace with the maximum of the left subtree
Return the root of the subtree
Must be recursive
Balances the tree if it needs it
O(log(n)) time complexity, O(1)* space complexity
:param node: root of the tree
:param value: value to be removed
:return: root of the subtree
"""
if node is not None:
if node.value < value:
# remove at right
self.remove(node.right, value)
elif node.value > value:
self.remove(node.left, value)
else: # find the element
parent = node.parent
if node.left is not None and node.right is not None:
if self.height(node.left) >= self.height(node.right):
node.value = self.max(node.left).value
self.remove(node.left, node.value)
else:
node.value = self.min(node.right).value
self.remove(node.right, node.value)
elif node.left is None and node.right is None:
if parent is None:
self.root = None
self.size = 0
elif parent.left == node:
parent.left = None
else:
parent.right = None
self.size -= 1
else:
if parent is None:
self.root = node = node.left if node.right is None else node.right
else:
flag = (node.parent.left == node)
if flag:
parent.left = node
else:
parent.right = node
node.parent = parent
self.size -= 1
node.height = max(self.height(node.left),
self.height(node.right)) + 1
self.rebalance(node)
def search(self, node, value):
"""
Takes in a value to search for and a node which is the root of a given tree or subtree
Returns the node with the given value if found, else returns the potential parent node
O(log(n)) time complexity, O(1)* space complexity
:param node: root of a given tree or subtree
:param value: value that is to be searched
:return: node with the given value or the potential parent node
"""
result = None
if node is not None:
if value < node.value:
result = self.search(node.left, value)
elif value > node.value:
result = self.search(node.right, value)
return result or node
def inorder(self, node):
"""
Takes in a value to search for and a node which is the root of a given tree or subtree
Returns the node with the given value if found, else returns the potential parent node
O(log(n)) time complexity, O(1)* space complexity
:param node: root of given tree or subtree
:return: list of the inorder result
"""
if node is not None:
yield from self.inorder(node.left)
yield node
yield from self.inorder(node.right)
def preorder(self, node):
"""
Same as inorder, only using the preorder method of traversal
O(n) time complexity, O(n) space complexity
:param node: root of a given tree or subtree
:return: list of the preorder result
"""
if node is not None:
yield node
yield from self.preorder(node.left)
yield from self.preorder(node.right)
def postorder(self, node):
"""
Same as inorder, only using the postorder method of traversal
O(n) time complexity, O(n) space complexity
:param node: root of a given tree or subtree
:return: list of the postorder result
"""
if node is not None:
yield from self.postorder(node.left)
yield from self.postorder(node.right)
yield node
def depth(self, value):
"""
Returns the depth of the node with the given value
O(height) time complexity, O(1) space complexity
:param value: the given value
:return: the depth of the node
"""
if self.root:
node = self.search(self.root, value)
if node.value == value:
depth = 0
while node != self.root:
node = node.parent
depth += 1
return depth
return -1
def height(self, node):
"""
Returns the height of the tree rooted at the given node
O(1) time complexity, O(1) space complexity
:param node: root of a given node
:return: the height of the node
"""
if node is not None:
return node.height
if self.root is None:
return 0
else:
return -1
def min(self, node):
"""
Returns the minimum of the tree rooted at the given node
Must be recursive
O(log(n)) time complexity, O(1)* space complexity
:param node: root of a given node
:return: minimum node of the tree
"""
if node is not None:
return self.min(node.left) or node
def max(self, node):
"""
Returns the maximum of the tree rooted at the given node
Must be recursive
O(log(n)) time complexity, O(1)* space complexity
:param node: root of a given tree
:return: maximum node of the tree
"""
if node is not None:
return self.max(node.right) or node
def get_size(self):
"""
Returns the number of nodes in the AVL Tree
O(1) time complexity, O(1) space complexity
:return: size of the AVL tree
"""
return self.size
def get_balance(self, node):
"""
Returns the balance factor of the node passed in
Balance Factor = height of left subtree – height of right subtree
O(1) time complexity, O(1) space complexity
:param node: a node in tree
:return: balance factor of the node
"""
if node is not None:
return self.height(node.left) - self.height(node.right)
else:
return 0
def left_rotate(self, root):
"""
Performs an AVL left rotation on the subtree rooted at root
Returns the root of the new subtree
O(1) time complexity, O(1) space complexity
:param root: root node of subtree
:return: root of the new subtree after left rotate
"""
new_root = root.right
root.right = new_root.left
if new_root.left:
new_root.left.parent = root
new_root.left = root
new_root.parent = root.parent
if root.parent:
if root.parent.left == root:
root.parent.left = new_root
else:
root.parent.right = new_root
root.parent = new_root
root.height = max(self.height(root.left), self.height(root.right)) + 1
new_root.height = max(self.height(new_root.right), root.height) + 1
p = new_root.parent
while p is not None:
p.height = max(self.height(p.left),
self.height(p.right)) + 1
p = p.parent
if root == self.root:
self.root = new_root
def right_rotate(self, root):
"""
Performs an AVL right rotation on the subtree rooted at root
Returns the root of the new subtree
O(1) time complexity, O(1) space complexity
:param root: root of the subtree
:return: new root of the subtree after right rotation
"""
new_root = root.left
root.left = new_root.right
if new_root.right:
new_root.right.parent = root
new_root.right = root
new_root.parent = root.parent
if root.parent:
if root.parent.left == root:
root.parent.left = new_root
else:
root.parent.right = new_root
root.parent = new_root
root.height = max(self.height(root.left), self.height(root.right)) + 1
new_root.height = max(self.height(new_root.left), root.height) + 1
p = new_root.parent
while p is not None:
p.height = max(self.height(p.left),
self.height(p.right)) + 1
p = p.parent
if root == self.root:
self.root = new_root
def rebalance(self, node):
"""
Rebalances the subtree rooted at node, if needed
Returns the root of the new, balanced subtree
O(1) time complexity, O(1) space complexity
:param node: root of subtree
:return: new root of subtree
"""
if node:
self.rebalance(node.left)
self.rebalance(node.right)
while abs(self.get_balance(node)) >= 2:
if self.get_balance(node) > 0:
if self.get_balance(node.left) >= 0:
self.right_rotate(node)
else:
self.left_rotate(node.left)
self.right_rotate(node)
else:
if self.get_balance(node.right) <= 0:
self.left_rotate(node)
else:
self.right_rotate(node.right)
self.left_rotate(node)
def repair_tree(tree):
"""
Takes in a tree where two values may have been swapped,
violating the BST property of nodes on the left being less
than the parent node, and nodes on the right being larger
than the parent node
Repairs the tree by finding if two values have actually been swapped,
and swapping them back if necessary
O(n) time complexity, O(n) space complexity
:param tree: a tree that may have two values swapped
:return: repaired tree
"""
if tree.root is None:
return tree
a = b = c = None
d = 0
for root in tree.inorder(tree.root):
if c and root.value < c.value:
if not a:
a, b = c, root
else:
b = root
d += 1
c = root
if d <= 2 and a and b:
a.value, b.value = b.value, a.value
return tree