Binary tree algorithms

binary_trees.py

#!/usr/bin/env python

from itertools import cycle

class Node(object):
	def __init__(self, value, left=None, right=None):
		self.value = value
		self.left = left
		self.right = right
	def preorder(self, shape=False):
		if self.left is None and self.right is None:
			return [0] if shape else [self.value]
		else:
			traversal = [1] if shape else [self.value]
			if self.left is not None:
				traversal.extend(self.left.preorder(shape))
			if self.right is not None:
				traversal.extend(self.right.preorder(shape))
			return traversal
	def mirror(self):
		if self.left is self.right is None:
			return
		temp = self.left
		self.left = self.right
		self.right = temp
		self.left.mirror()
		self.right.mirror()


class BinaryTree(object):
	def __init__(self, root):
		self.root = root

	def preorder(self):
		# preorder traversal, emit: 0 if leaf, 1 if internal
		traversal = []
		traversal.extend(self.root.preorder())
	def mirror(self):
		self.root.mirror()
	def __eq__(self, other):
		return self.root.preorder(shape=True) == other.root.preorder(shape=True)
	def __hash__(self):
		return str(self.root.preorder())
	def __repr__(self):
		return str(self.root.preorder())


def build_tree(n):
	# n is a name generator
	# build full binary tree of depth 2
	btree = BinaryTree(Node(next(n)))
	btree.root.left = Node(next(n))
	btree.root.left.left = Node(next(n))
	btree.root.left.right = Node(next(n))
	btree.root.right = Node(next(n))
	btree.root.right.left = Node(next(n))
	btree.root.right.right = Node(next(n))
	return btree

def build_tree2(n):
	# n is a name generator
	# build incomplete binary tree of depth 3
	btree = BinaryTree(Node(next(n)))
	btree.root.left = Node(next(n))
	btree.root.left.left = Node(next(n))
	btree.root.left.right = Node(next(n))
	btree.root.left.left.left = Node(next(n))
	btree.root.left.left.right = Node(next(n))
	btree.root.right = Node(next(n))
	btree.root.right.left = Node(next(n))
	btree.root.right.right = Node(next(n))
	return btree

def build_tree3(n):
	# n is a name generator
	# build incomplete binary tree of depth 3
	btree = BinaryTree(Node(next(n)))
	btree.root.left = Node(next(n))
	btree.root.left.left = Node(next(n))
	btree.root.left.right = Node(next(n))
	btree.root.right = Node(next(n))
	btree.root.right.left = Node(next(n))
	btree.root.right.right = Node(next(n))
	btree.root.right.right.left = Node(next(n))
	btree.root.right.right.right = Node(next(n))
	return btree

def compare_trees(bt1, bt2):
	print "bt1: {}".format(bt1)
	print "bt2: {}".format(bt2)
	print "bt1 and bt2 are "  + "equal" if bt1 == bt2 else "different"

def node_letter():
	for letter in cycle(xrange(ord('A'), ord('Z'))):
		yield chr(letter)

def main():
	n = node_letter()
	bt1 = build_tree(n)
	bt2 = build_tree(n)
	compare_trees(bt1, bt2)

	bt3 = build_tree2(n)
	bt4 = build_tree3(n)
	compare_trees(bt3, bt4)

	bt3.mirror()
	compare_trees(bt3, bt4)

if __name__=="__main__":
	main()