Tree.py

# Trees

def tree(root_label, branches = []):
    for branch in branches:
        assert is_tree(branch),'branches must be trees'
    return [root_label] + list(branches)


def label(tree):
    return tree[0]

def branches(tree):
    return tree[1:]

def is_tree(tree):
    if type(tree) != list or len(tree) < 1:
        return False
    for branch in branches(tree):
        if not is_tree(branch):
            return False
    return True

def is_leaf(tree):
    return not branches(tree)

def fib_tree(n):
    if n <= 1:
        return tree(n)     # here is a leaf, but must construct as a tree
    else:
        left, right = fib_tree(n - 2), fib_tree(n - 1)
        return tree(label(left)+label(right), [left, right])
    
def count_leaf(t):
    """Count the leaves of a tree."""
    if is_leaf(t):
        return 1
    else:
        branch_counts = [count_leaf(b) for b in branches(t)]  
        return sum(branch_counts)
    
def leaves(tree):
    """Return  a list containing the leaf labels of tree.
    
    >>> leaves(fib_tree(5))
    [1, 0, 1, 0, 1, 1, 0, 1]
    """
    if is_leaf(tree):
        return [label(tree)]
    else:
        return sum([leaves(b) for b in branches(tree)], [])

def increment_leaves(t):
    """Return a tree like t but with leaf labels incremented.

    
    """
    if is_leaf(t):
        return tree(label(t) + 1)
    else:
        bs = [increment_leaves(b) for b in branches(t)]
        return tree(label(t), bs)
    

def increment(t):
    """Return a tree like t but with all labels incremented."""
    
    return tree(label(t) + 1, [increment(b) for b in branches(t)])

def print_tree(t, indent = 0):
    if is_leaf(t):
        print(' ' * indent + str(label(t)))
    else:
        print(' ' * indent + str(label(t)))
        for b in branches(t):
            print_tree(b, indent + 1)