left_leaning_red_black_tree.py

#LLRBT implementation
from binary_search_tree import BSTNode

class LLRBNode(BSTNode):
  def __init__(self, black: bool, key, val, left = None, right = None, parent = None) -> None:
    self.black = black
    self.key = key
    self.val = val
    self.left = left
    self.right = right
    self.parent = parent

  def set_black(self, black: bool):
    self.black = black
  def flip_color(self):
    self.black = not self.black

class LLRB:
  def __init__(self):
    self.__root = None
  
  def insert(self, key, val):
    self.__root = self.__insert(root, val)
    self.__root.set_black(True)

  def __insert(self, node: LLRBNode, key, val) -> LLRBNode:
    # Check whether this is the first ever node inserted
    if node is None:
      return LLRBNode(True, key, val)
    # Normal BST search logic
    if key == node.key:
      return node
    elif key > node.key:
      node.right = self.__insert(node.right, key, val)
    else:
      node.left = self.__insert(node.left, key, val)
    # 3 fixes, in this order
    if LLRB.__is_red(node.right) and not LLRB.__is_red(node.left):
      node = LLRB.__rotate_left(node)
    if LLRB.__is_red(node.left) and LLRB.__is_red(node.left.left):
      node = LLRB.__rotate_right(node)
    if LLRB.__is_red(node.left) and LLRB.__is_red(node.right):
      node = LLRB.__flip_color(node)
    return node
  
  def find(self, key):
    return self.__root.find(key)

  @staticmethod
  def __is_red(node:LLRBNode) -> bool:
    return node and not node.black
  @staticmethod
  def __rotate_right(node: LLRBNode) -> LLRBNode:
    new_root = node.left
    node.left, new_root.right = new_root.right, node
    new_root.set_black(node.set_black)
    node.set_black(False)
    return new_root
  @staticmethod
  def __rotate_left(node: LLRBNode) -> LLRBNode:
    new_root = node.right
    node.right, new_root.left = new_root.left, node
    new_root.set_black(node.set_black)
    node.set_black(False)
    return new_root
  @staticmethod
  def __flip_color(node: LLRBNode) -> None:
    node.flip_color()
    node.left.flip_color()
    node.right.flip_color()