class Tree(object): """A Tree consists of a label and a sequence of 0 or more Trees, called its children.""" def __init__(self, entry, *children): """A Tree with given entry and children. For convenience, if children[k] is not a Tree, it is converted into a leaf whose operator is children[k]. """ self.entry = entry; self.children = [c if type(c) is Tree else Tree(c) for c in children] @property def is_leaf(self): """A tree is a leaf if it does not have any children.""" return len(self.children) == 0 @property def items(self): """Returns a tuple of all items in a given tree. >>> t = Tree(1, Tree(2, Tree(3), Tree(4), Tree(5))) >>> t.items (1, 2, 3, 4, 5) """ tree_items = (self.entry,) for child in self.children: tree_items += child.items return tree_items @property def size(self): """Returns the total number of items in a tree. >>> t = Tree(1, Tree(2, Tree(3)), Tree(4), Tree(5)) >>> t.size 5 """ return len(self.items) def copy(self): """Create a copy of the tree. >>> t = Tree(2, Tree(3, Tree(4), Tree(5), Tree(7))) >>> s = t.copy() >>> s.entry = 6 >>> s Tree(6, Tree(3, Tree(4), Tree(5), Tree(7))) >>> t Tree(2, Tree(3, Tree(4), Tree(5), Tree(7))) """ new_tree = Tree(self.entry) if self.is_leaf: return new_tree for child in self.children: new_tree.children.append(child.copy()) return new_tree def map(self, fn): """Applies the function, fn, to each element of the tree and mutates the original tree. Does not create a new tree. >>> t = Tree(1, Tree(2), Tree(3), Tree(4)) >>> t.map(lambda x: x * x) >>> t Tree(1, Tree(4), Tree(9), Tree(16)) """ self.entry = fn(self.entry) for child in self.children: child.map(fn) def __str__(self): return repr(self) def __getitem__(self, k): """Gets the child at index k. >>> t = Tree(4, Tree(5), Tree(6), Tree(1)) >>> t[1] Tree(6) """ return self.children[k] def __iter__(self): return iter(self.children) def __repr__(self): if self.is_leaf: return "Tree({0})".format(self.entry) else: return "Tree({0}, {1})".format(self.entry, str(self.children)[1:-1])