binary_tree.py

#Binary Tree implementation
from collections import deque

class BinaryTreeNode:
  def __init__(self, label, left = None, right = None) -> None:
    self.left = left
    self.right = right
    self.label = label

  def __str__(self) -> str:
    return self.label

  def show(self):
    self.__indent_print("", "")

  def __indent_print(self, prefix, branchName):
    print(prefix + branchName + self.__str__())
    if self.has_left:
      self.left.__indent_print(BinaryTreeNode.__next_prefix(prefix), "(LEFT)")
    if self.has_right:
      self.right.__indent_print(BinaryTreeNode.__next_prefix(prefix), "(RIGHT)")
  
  @staticmethod
  def __next_prefix(str) -> str:
    if str == '':
      return '├─ '
    else:
      return '│  ' + str

  @property
  def has_left(self) -> bool:
    return self.left is not None

  @property
  def has_right(self) -> bool:
    return self.right is not None
  
  @property
  def arity(self) -> int:
    count = 0
    if self.has_left:
      count += 1
    if self.has_right:
      count += 1
    return count

  @property
  def is_leaf(self) -> bool:
    return self.arity == 0

  def pre_order(self, action) -> None:
    if self is None: # or, if self.
      # I like to keep it straightforward. Sometimes shorthand
      # confuses between empty collection and real None
      return
    action(self.label)
    if self.has_left:
      self.left.pre_order(action)
    if self.has_right:
      self.right.pre_order(action)
  # If the children are stored in a list, just traverse through the list

  def in_order(self, action) -> None:
    if self is None:
      return
    if self.has_left:
      self.left.in_order(action)
    action(self.label)
    if self.has_right:
      self.right.in_order(action)

  def post_order(self, action) -> None:
    if self is None:
      return
    if self.has_left:
      self.left.post_order(action)
    if self.has_right:
      self.right.post_order(action)
    action(self.label)

  # We can make the above methods static to avoid "None check"s on children.
  # This is somewhat cleaner, saving two lines
  # For example:
  # @staticmethod
  # def spost_order(node, action) -> None:
  #   if node is None:
  #     return
  #   BinaryTreeNode.spost_order(node.left, action)
  #   BinaryTreeNode.spost_order(node.right, action)
  #   action(node.label)

  # bfs
  def level_order(self, action) -> None:
    if self is None:
      return
    # deque is similar to ArrayDeque in Java
    # we can use a vanilla list here as well: list = [self]
    queue = deque()
    queue.append(self)
    while queue:
      cur = queue.popleft() # same as list.pop(0), which is slower
      action(cur)
      # It is important here to add children in order
      if cur.hasLeft:
        queue.append(cur.left)
      if cur.hasRight:
        queue.append(cur.right)
    return
  
  # an iterative version of the pre-order dfs
  def pre_order_iterative(self, action) -> None:
    if self is None:
      return
    stack = deque()
    stack.append(self)
    while stack:
      cur = stack.pop()
      action(cur)
      # It is important to add children in reversed order 
      # This is to ensure pre-order traversal
      if cur.hasRight:
        stack.append(cur.right)
      if cur.hasLeft:
        stack.append(cur.left)
  
  # an iterative version of the in-order dfs
  # somewhat more complex than iterative pre-order
  def in_order_iterative(self, action) -> None:
    cur, stack = self, deque()
    while True:
      if cur is not None:
        stack.append(cur)
        cur = cur.left
      elif stack:
        cur = stack.pop()
        action(cur)
        cur = cur.right
      else:
        break

  # an iterative version of the post-order dfs
  # it will be more complex than the previous two due to its non-tail recursive nature
  # we can cheat our way through using two stacks by using a 
  # strategy somewhat like pre-order iterative
  def post_order_iterative_two_stack(self, action) -> None:
    if self is None:
      return
    stack, stack2 = deque(), deque()
    stack.append(self)
    while stack:
      cur = stack.pop()
      stack2.append(cur)
      # Here, we add children in correct order
      if cur.hasLeft:
        stack.append(cur.left)
      if cur.hasRight:
        stack.append(cur.right)
    while stack2:
      action(stack2.pop())
  
  # it is possible with one stack
  # this strategy is somewhat like in-order iterative
  def post_order_iterative_one_stack(self, action) -> None:
    cur, stack = self, deque()
    while True:
      while cur:
        stack.append(cur)
        stack.append(cur) # add root node twice
        cur = cur.left
      if not stack:
        break
      cur = stack.pop()
      if stack and stack.peek() == cur:
        cur = cur.right
      else:
        action(cur)
        cur = None
    
  # Iterative deepening:
  # somewhat a combination of dfs and bfs
  # do breadth-first traversal by repeated (truncated) depthfirst traversals
  def iterative_deepening(self, action, max_depth: int) -> None:
    if self is None or max_depth < 0:
      return
    for i in range(0, max_depth):
      self.__do_level(action, i)
  
  def __do_level(self, action, level):
    if level == 0:
      action(self)
    else:
      if self.hasLeft:
        self.__do_level(action, level - 1)
      if self.hasRight:
        self.__do_level(action, level - 1)