algo to clone a graph

clone_graph_2.py

 
class TreeNode:
  """ TreeNode has a value with 3 pointers """
  def __init__(self, value, left=None, right=None, random=None):
    self.value = value
    self.left = left
    self.right = right
    self.random = random
  
  def __str__(self):
    return "(%s)"%(self.value)
  
def build_a_tree():
  root = TreeNode(10)
  a = TreeNode(20)
  b = TreeNode(30)
  c = TreeNode(40)
  d = TreeNode(50)
  
  root.left = a
  root.right = b
  a.right = c
  b.left = d

  a.random = c
  d.random = b

  return root


def copy_left_right(root, d):
  """ Step 1: Copy the left & right pointers.
  root mapped to cloned in a dictionary """
  if root is None:
    return None
  
  cloned = TreeNode(root.value)
  d[root] = cloned
  cloned.left = copy_left_right(root.left, d)
  cloned.right = copy_left_right(root.right, d)
  return cloned

def copy_random(root, cloned, d):
  """ Step 2: Copy the random pointer """
  if cloned is None:
    return
  
  cloned.random = d.get(root.random)
  copy_random(root.left, cloned.left, d)
  copy_random(root.right, cloned.right, d)
  

def clone_tree(root):
  """clone a tree with random pointers"""
  if root is None:
    return None
  
  d = {}
  cloned = copy_left_right(root, d)
  copy_random(root, cloned, d)
  return cloned

"""driver program"""

root = build_a_tree()

new_root = clone_tree(root)

print("tree node values")
print(new_root)
print(new_root.left)
print(new_root.right)
print(new_root.left.right)
print(new_root.right.left)

print("random node values")
print(new_root.left.random) 
print(new_root.right.left.random)