RBTNode.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sun Nov  4 20:21:24 2018

@author: JesusMHernandez
"""

class RBTNode:
  def __init__(self, key, parent, is_red = False, left = None, right = None):
    self.key = key
    self.left = left
    self.right = right
    self.parent = parent
        
    if is_red:
      self.color = "red"
    else:
      self.color = "black"

  # Returns true if both child nodes are black. A child set to None is considered
  # to be black.
  def are_both_children_black(self):
    if self.left != None and self.left.is_red():
      return False
    if self.right != None and self.right.is_red():
      return False
    return True

  def count(self):
    count = 1
    if self.left != None:
      count = count + self.left.count()
    if self.right != None:
      count = count + self.right.count()
    return count
    
  # Returns the grandparent of this node
  def get_grandparent(self):
    if self.parent is None:
      return None
    return self.parent.parent

  # Gets this node's predecessor from the left child subtree
  # Precondition: This node's left child is not None
  def get_predecessor(self):
    node = self.left
    while node.right is not None:
      node = node.right
    return node

  # Returns this node's sibling, or None if this node does not have a sibling
  def get_sibling(self):
    if self.parent is not None:
      if self is self.parent.left:
        return self.parent.right
      return self.parent.left
    return None

  # Returns the uncle of this node
  def get_uncle(self):
    grandparent = self.get_grandparent()
    if grandparent is None:
      return None
    if grandparent.left is self.parent:
      return grandparent.right
    return grandparent.left

  # Returns True if this node is black, False otherwise
  def is_black(self):
    return self.color == "black"

  # Returns True if this node is red, False otherwise
  def is_red(self):
    return self.color == "red"

  # Replaces one of this node's children with a new child
  def replace_child(self, current_child, new_child):
    if self.left is current_child:
      return self.set_child("left", new_child)
    elif self.right is current_child:
      return self.set_child("right", new_child)
    return False

  # Sets either the left or right child of this node
  def set_child(self, which_child, child):
    if which_child != "left" and which_child != "right":
      return False

    if which_child == "left":
      self.left = child
    else:
      self.right = child

    if child != None:
      child.parent = self

    return True