build binary tree

buildBinaryTree.py

# Question: to build a binary tree, with specific input format:


# input: 
# 2 3 4 (2, 3, 4 is the index of a node)
# 3 7 8
# 8 1 a

#    2
#  /  \
# 3   4
# /\
# 7 8
#  / \
# 1  a

# output:
# post-order travesal of the tree...
# 7 1 a 8 3 4 2 

# My Answer:
# step 1. define data structure, how to store the data
class Node:
    def __init__(self, idx):
        self.idx = idx
        self.left = None
        self.right = None
        
    def __repr__(self):
        return "TreeNode: %d" % self.idx
    
    @staticmethod
    def connect(this, left, right):
        this.left = left
        this.right = right
        
        return this
        
    
# step 2. define input as python object
rows = [
    [2, 3, 4],
    [3, 7, 8],
    [8, 1, 0]  # guarantee the input is valid.. 
]
        
    
# step 3. define function take input as parameter, and consider what should return as representation of the tree
def find_or_create_node(created_nodes, idx):
    node = created_nodes[idx]
    if node is None:
        # not created yet, create one and populate back
        node = Node(idx)
        created_nodes[idx] = node
    return node
    
def build_tree(rows):
    created_nodes = [None] * 999
    node_occurence_cnt = {}
    for row in rows:
        nodes = [None, None, None]
        
        # if it's parent node, add one; otherwise decrease one.
        # result is only root node got number 1
        node_occurence_cnt[row[0]] = node_occurence_cnt.get(row[0], 0) + 1
        node_occurence_cnt[row[1]] = node_occurence_cnt.get(row[1], 0) - 1
        node_occurence_cnt[row[2]] = node_occurence_cnt.get(row[2], 0) - 1
        
        for idx in range(3):
            nodes[idx] = find_or_create_node(created_nodes, row[idx])
        Node.connect(nodes[0], nodes[1], nodes[2])

    # find root node, which is the one doesn't have parent.
    # after the above operations, there should be only one root left.
    for idx, cnt in node_occurence_cnt.items():
        # only appears once as parent
        if cnt == 1:
            return created_nodes[idx]

def walk(root):
    if root.left:
        walk(root.left)
    if root.right:
        walk(root.right)
    print root
    
    
root = build_tree(rows)
walk(root)