Demonstrating depth-first search and breadth-first-search in Python

trees.py

import random


class TreeNode(object):
    def __init__(self, value, left=None, right=None):
        self.value = value
        self.left = left
        self.right = right

    def __hash__(self):
        return hash(self.value)

    def __repr__(self):
        return "TreeNode(value=%r, left=%r, right=%r)" % (self.value, self.left, self.right)

    def __str__(self):
        return self._str(0)

    def _str(self, indent):
        def _str_children(child):
            if child is not None:
                return child._str(indent + 2)
            else:
                return ' ' * (indent + 2) + '-'
        parts = [
            ' ' * indent + str(self.value),
            _str_children(self.left),
            _str_children(self.right),
        ]
        return "\n".join(parts)


def insert(root, node):
    if node.value <= root.value:
        if root.left is None:
            root.left = node
        else:
            insert(root.left, node)
    else:
        if root.right is None:
            root.right = node
        else:
            insert(root.right, node)


def dfs(node, visited=None):
    if visited is None:
        visited = set()
    print node.value
    visited.add(node)
    visit = lambda n: n is not None and n not in visited
    if visit(node.left):
        dfs(node.left, visited)
    if visit(node.right):
        dfs(node.right, visited)


# vals = range(1, 13)
# random.shuffle(vals)
vals = [4, 10, 9, 5, 2, 7, 3, 1, 6, 12, 8, 11]
print vals
print
root = TreeNode(vals[0])
for val in vals[1:]:
    insert(root, TreeNode(val))
print root
print

dfs(root)