convert.py

class Tree:

    def __init__(self, base, left, right):
        self.base = base
        self.left = left
        self.right = right

    def __repr__(self):
        return "(%s, %s, %s)" % (self.base, self.left, self.right)


def inorder(tree):
    if tree is None:
        return
    for el in inorder(tree.left):
        yield el
    yield tree.base
    for el in inorder(tree.right):
        yield el

def preorder(tree):
    if tree is None:
        return
    yield tree.base
    for el in inorder(tree.left):
        yield el
    for el in inorder(tree.right):
        yield el

def restore(inorder_seq, preorder_seq):
    assert len(inorder_seq) == len(preorder_seq)
    if len(preorder_seq) == 0:
        return None
    base = preorder_seq[0]
    base_idx = inorder_seq.index(base)
    left_inorder = inorder_seq[:base_idx]
    right_inorder = inorder_seq[base_idx+1:]
    left_preorder = preorder_seq[1:][:base_idx]
    right_preorder = preorder_seq[1:][base_idx:]
    return Tree(base, restore(left_inorder, left_preorder), restore(right_inorder, right_preorder))


def test_tree(tree):
    inorder_seq = list(inorder(tree))
    preorder_seq = list(preorder(tree))
    restored_tree = restore(inorder_seq, preorder_seq)
    assert inorder_seq == list(inorder(restored_tree))
    assert preorder_seq == list(preorder(restored_tree))

t = Tree(1, Tree(2, None, None), Tree(3, None, None))
print list(inorder(t))
print list(preorder(t))

test_tree(Tree(1, Tree(2, Tree(3, None, None), None), None))
test_tree(Tree(1, Tree(2, None, None), Tree(3, None, None)))
test_tree(Tree(1, Tree(2, None, None), None))
test_tree(Tree(1, None, Tree(2, None, None)))
test_tree(Tree(1, Tree(2, None, None), Tree(3, None, None)))
test_tree(Tree(1, Tree(2, None, None), None))