288 lines
9.2 KiB
Python
288 lines
9.2 KiB
Python
from dataclasses import dataclass
|
|
from enum import Enum
|
|
from tools.lists import Queue
|
|
from typing import Any, Union
|
|
|
|
|
|
class Rotate(Enum):
|
|
LEFT = 0
|
|
RIGHT = 1
|
|
|
|
|
|
@dataclass
|
|
class TreeNode:
|
|
value: Any
|
|
parent: Union['TreeNode', None] = None
|
|
left: Union['TreeNode', None] = None
|
|
right: Union['TreeNode', None] = None
|
|
balance_factor: int = 0
|
|
height: int = 0
|
|
|
|
def __str__(self):
|
|
return "TreeNode:(%s; bf: %d, d: %d, p: %s, l: %s, r: %s)" \
|
|
% (self.value, self.balance_factor, self.height,
|
|
self.parent.value if self.parent else "None",
|
|
self.left.value if self.left else "None",
|
|
self.right.value if self.right else "None")
|
|
|
|
def __repr__(self):
|
|
return str(self)
|
|
|
|
|
|
def update_node(node: TreeNode):
|
|
left_depth = node.left.height if node.left is not None else -1
|
|
right_depth = node.right.height if node.right is not None else -1
|
|
node.height = 1 + max(left_depth, right_depth)
|
|
node.balance_factor = right_depth - left_depth
|
|
|
|
|
|
class BinarySearchTree:
|
|
root: Union[TreeNode, None] = None
|
|
node_count: int = 0
|
|
|
|
def _balance(self, node: TreeNode) -> TreeNode:
|
|
if node.balance_factor == -2:
|
|
if node.left.balance_factor <= 0:
|
|
return self.rotate(Rotate.RIGHT, node)
|
|
else:
|
|
return self.rotate(Rotate.RIGHT, self.rotate(Rotate.LEFT, node.left))
|
|
elif node.balance_factor == 2:
|
|
if node.right.balance_factor >= 0:
|
|
return self.rotate(Rotate.LEFT, node)
|
|
else:
|
|
return self.rotate(Rotate.LEFT, self.rotate(Rotate.RIGHT, node.right))
|
|
else:
|
|
return node
|
|
|
|
def _insert(self, node: TreeNode, parent: TreeNode, obj: Any) -> TreeNode:
|
|
if node is None:
|
|
return TreeNode(obj, parent)
|
|
|
|
if obj < node.value:
|
|
node.left = self._insert(node.left, node, obj)
|
|
else:
|
|
node.right = self._insert(node.right, node, obj)
|
|
|
|
update_node(node)
|
|
return self._balance(node)
|
|
|
|
def add(self, obj: Any):
|
|
if obj is None or obj in self:
|
|
raise ValueError("obj is None or already present in tree")
|
|
|
|
self.root = self._insert(self.root, self.root, obj)
|
|
self.node_count += 1
|
|
|
|
def remove(self, obj: Any, root_node: TreeNode = None):
|
|
if self.root is None:
|
|
raise IndexError("remove from empty tree")
|
|
if root_node is None:
|
|
root_node = self.root
|
|
|
|
node = root_node
|
|
while node is not None:
|
|
if obj < node.value:
|
|
node = node.left
|
|
continue
|
|
elif obj > node.value:
|
|
node = node.right
|
|
continue
|
|
else:
|
|
if node.left is None and node.right is None: # leaf node
|
|
if node.parent is not None:
|
|
if node.parent.left == node:
|
|
node.parent.left = None
|
|
else:
|
|
node.parent.right = None
|
|
elif node.left is not None and node.right is not None: # both subtrees present
|
|
d_node = node.left
|
|
while d_node.right is not None:
|
|
d_node = d_node.right
|
|
node.value = d_node.value
|
|
self.remove(node.value, d_node)
|
|
elif node.left is None: # only a subtree on the right
|
|
if node.parent is not None:
|
|
if node.parent.left == node:
|
|
node.parent.left = node.right
|
|
else:
|
|
node.parent.right = node.right
|
|
node.right.parent = node.parent
|
|
else: # only a subtree on the left
|
|
if node.parent is not None:
|
|
if node.parent.left == node:
|
|
node.parent.left = node.left
|
|
else:
|
|
node.parent.right = node.left
|
|
node.left.parent = node.parent
|
|
|
|
update_node(root_node)
|
|
self._balance(root_node)
|
|
return
|
|
|
|
raise ValueError("obj not in tree:", obj)
|
|
|
|
def rotate(self, direction: Rotate, node: TreeNode = None) -> TreeNode:
|
|
if node is None:
|
|
node = self.root
|
|
|
|
parent = node.parent
|
|
if direction == Rotate.LEFT:
|
|
pivot = node.right
|
|
node.right = pivot.left
|
|
if pivot.left is not None:
|
|
pivot.left.parent = node
|
|
pivot.left = node
|
|
else:
|
|
pivot = node.left
|
|
node.left = pivot.right
|
|
if pivot.right is not None:
|
|
pivot.right.parent = node
|
|
pivot.right = node
|
|
|
|
node.parent = pivot
|
|
pivot.parent = parent
|
|
|
|
if parent is not None:
|
|
if parent.left == node:
|
|
parent.left = pivot
|
|
else:
|
|
parent.right = pivot
|
|
|
|
if node == self.root:
|
|
self.root = pivot
|
|
|
|
update_node(node)
|
|
update_node(pivot)
|
|
|
|
return pivot
|
|
|
|
def print(self, node: TreeNode = None, level: int = 0):
|
|
if node is None:
|
|
if level == 0 and self.root is not None:
|
|
node = self.root
|
|
else:
|
|
return
|
|
|
|
self.print(node.right, level + 1)
|
|
print(" " * 4 * level + '->', node)
|
|
self.print(node.left, level + 1)
|
|
|
|
def __contains__(self, obj: Any) -> bool:
|
|
if self.root is None:
|
|
return False
|
|
|
|
c_node = self.root
|
|
while c_node is not None:
|
|
if obj == c_node.value:
|
|
return True
|
|
elif obj < c_node.value:
|
|
c_node = c_node.left
|
|
else:
|
|
c_node = c_node.right
|
|
|
|
return False
|
|
|
|
def __len__(self) -> int:
|
|
return self.node_count
|
|
|
|
|
|
class Heap(BinarySearchTree):
|
|
def _find_left(self) -> TreeNode:
|
|
heap = Queue()
|
|
heap.put(self.root)
|
|
while c_node := heap.get():
|
|
if c_node.left is None or c_node.right is None:
|
|
return c_node
|
|
else:
|
|
heap.put(c_node.left)
|
|
heap.put(c_node.right)
|
|
|
|
def _find_right(self) -> TreeNode:
|
|
heap = [self.root]
|
|
while c_node := heap.pop():
|
|
if c_node.left is None and c_node.right is None:
|
|
return c_node
|
|
else:
|
|
if c_node.left is not None:
|
|
heap.append(c_node.left)
|
|
if c_node.right is not None:
|
|
heap.append(c_node.right)
|
|
|
|
def _sort_up(self, node: TreeNode):
|
|
pass
|
|
|
|
def _heapify(self):
|
|
pass
|
|
|
|
def empty(self) -> bool:
|
|
return self.root is None
|
|
|
|
def add(self, obj: Any):
|
|
node = TreeNode(obj)
|
|
if self.root is None:
|
|
self.root = node
|
|
else:
|
|
t_node = self._find_left()
|
|
if t_node.left is None:
|
|
t_node.left = node
|
|
else:
|
|
t_node.right = node
|
|
node.parent = t_node
|
|
self._sort_up(node)
|
|
|
|
def pop(self) -> Any:
|
|
if self.root is None:
|
|
raise IndexError("pop from empty heap")
|
|
|
|
ret = self.root.value
|
|
if self.root.left is None and self.root.right is None:
|
|
self.root = None
|
|
else:
|
|
d_node = self._find_right()
|
|
self.root.value = d_node.value
|
|
if d_node.parent.left == d_node:
|
|
d_node.parent.left = None
|
|
else:
|
|
d_node.parent.right = None
|
|
self._heapify()
|
|
|
|
return ret
|
|
|
|
|
|
class MinHeap(Heap):
|
|
def _sort_up(self, node: TreeNode):
|
|
while node.parent is not None and node.value < node.parent.value:
|
|
node.value, node.parent.value = node.parent.value, node.value
|
|
node = node.parent
|
|
|
|
def _heapify(self):
|
|
node = self.root
|
|
while node is not None:
|
|
if node.left and node.left.value < node.value and (not node.right or node.right.value > node.left.value):
|
|
node.left.value, node.value = node.value, node.left.value
|
|
node = node.left
|
|
elif node.right and node.right.value < node.value and (not node.left or node.left.value > node.right.value):
|
|
node.right.value, node.value = node.value, node.right.value
|
|
node = node.right
|
|
else:
|
|
break
|
|
|
|
|
|
class MaxHeap(Heap):
|
|
def _sort_up(self, node: TreeNode):
|
|
while node.parent is not None and node.value > node.parent.value:
|
|
node.value, node.parent.value = node.parent.value, node.value
|
|
node = node.parent
|
|
|
|
def _heapify(self):
|
|
node = self.root
|
|
while node is not None:
|
|
if node.left and node.left.value > node.value and (not node.right or node.right.value < node.left.value):
|
|
node.left.value, node.value = node.value, node.left.value
|
|
node = node.left
|
|
elif node.right and node.right.value > node.value and (not node.left or node.left.value < node.right.value):
|
|
node.right.value, node.value = node.value, node.right.value
|
|
node = node.right
|
|
else:
|
|
break
|