#!/usr/bin/env python


class Node:
    def __init__(self, val):
        self.right = None
        self.left = None
        self.val = val


def tree_insert(root, val):
    if val > root.val:
        if root.right is None:
            root.right = Node(val)
            return
        tree_insert(root.right, val)
    else:
        if root.left is None:
            root.left = Node(val)
            return
        tree_insert(root.left, val)


def tree_print(root):
    def _tree_print(root):
        if root is None:
            return
        _tree_print(root.left)
        if root.left:
            print("    {0} -> {1} [color=blue];".format(root.val, root.left.val))
        if root.right:
            print("    {0} -> {1} [color=red];".format(root.val, root.right.val))
        _tree_print(root.right)

    print("digraph G {")
    _tree_print(root)
    print("}")


# This is not optimal, the version below is better
def bstget(root, n):
    stack = []
    visited = []
    i = 0
    it = root

    while it:
        if it in visited:
            print(it.val)
            i += 1
            if i == n:
                return it
        else:
            if it.right:
                stack.append(it.right)
            stack.append(it)
            if it.left:
                stack.append(it.left)
            visited.append(it)
        it = stack.pop() if stack else None

    return None


def bstget2(root, n):
    stack = []
    i = 0
    it = root
    while stack or it:
        if it:
            stack.append(it)
            it = it.left
        else:
            it = stack.pop()
            print(it.val)
            i += 1
            if i == n:
                return it
            it = it.right

    return None

if __name__ == "__main__":
    root = Node(0)
    tree_insert(root, 2)
    tree_insert(root, 9)
    tree_insert(root, 1)
    tree_insert(root, 5)
    tree_insert(root, 6)
    tree_insert(root, 3)
    tree_insert(root, 4)
    tree_print(root)
    print("4th node is {0}".format(bstget(root, 4).val))
    print("4th node is {0}".format(bstget2(root, 4).val))