Wednesday, October 12, 2005

Tree Function examples

#trees.py
#
#See Roshan's blog
#

from itertools import izip

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

    def __str__(self):
        return "Node(%s, %s)"%(self.left, self.right)

def flatten(node):
    """
    Flatten the tree depth first from the left to right.
    """

    if isinstance(node, Node):
        for i in flatten(node.left):
            yield i
        for i in flatten(node.right):
            yield i
    else:
        yield node

def samefringe(tree1, tree2):
    """
    Check whether two trees have the same leaf nodes
    with the same value.
    """

    for leaf1, leaf2 in izip(flatten(tree1), flatten(tree2)):
        if leaf1!=leaf2:
            return False
    return True


def repmin(tree):
    """
    Return a new tree in with the same node structure
    but all leaf nodes have the same value which is the
    minimum leaf value.
    """

    minval, f = _repmin(tree)
    return f(minval)

def _repmin(tree):
    if isinstance(tree, Node):
       (minleft, left) = _repmin(tree.left)
       (minright, right) = _repmin(tree.right)
       return (min(minleft, minright),
               lambda minval: Node(left(minval), right(minval)))
    else:
        return (tree, lambda val: val)

tree1 = Node(1, Node(2, Node(3, Node(45))))
tree2 = Node(Node(1, Node(23)), Node(45))
tree3 = Node(Node(1, Node(24)), Node(45))

print samefringe(tree1, tree2)
print samefringe(tree1, tree3)

print repmin(tree1)

0 Comments:

Post a Comment

<< Home