binary_search_tree.py

class Node:
    def __init__(self, x, parent=None):
        self.key = x
        self.left = None
        self.right = None
        self.parent = parent

    def append_child(self, child):
        if self.key > child.key:
            self.left = child
            child.parent = self
        elif self.key < child.key:
            self.right = child
            child.parent = self
        else:
            raise "cannot append a child with same key"

    def accept(self, visitor):
        visitor.visit(self)

    def remove(self, tree):
        parent = self.parent
        if self.left and self.right:
            child = self.right
            while child.left is not None:
                child = child.left
            child.remove(tree)
            self.key = child.key
        elif self.left or self.right:
            child = self.left or self.right
            if parent is None: # if this is root
                tree.root = child
                child.parent = None
            else:
                parent.append_child(child)
        else:
            if parent is None: # if this is root
                tree.root = None
            elif self is parent.left:
                parent.left = None
            else:
                parent.right = None

class BinaryTree:
    def __init__(self):
        self.root = None

    def insert(self, x):
        inserted = Node(x)
        if self.root is None:
            self.root = inserted
            return inserted

        node = self.root
        parent = None
        while node is not None:
            if x > node.key:
                parent = node
                node = node.right
            elif x < node.key:
                parent = node
                node = node.left
            else:
                return node

        parent.append_child(inserted)
        return inserted

    def delete(self, x):
        deleted = self.find(x)
        if deleted is None:
            return False
        deleted.remove(self)
        return True

    def find(self, x):
        node = self.root
        while node is not None:
            if x == node.key:
                return node
            elif x > node.key:
                node = node.right
            else:
                node = node.left
        return None

    def pre_order(self):
        collector = PreOrderKeyCollector()
        if self.root is not None:
            self.root.accept(TreeVisitor(collector))
        print(collector.getResult())

    def in_order(self):
        collector = InOrderKeyCollector()
        if self.root is not None:
            self.root.accept(TreeVisitor(collector))
        print(collector.getResult())

    def max_level(self):
        counter = MaxLevelCounter()
        if self.root is not None:
            self.root.accept(TreeVisitor(counter))
        return counter.getResult()

    def count(self):
        counter = NodeCounter()
        if self.root is not None:
            self.root.accept(TreeVisitor(counter))
        return counter.getResult()

# yieldを使う方がスッキリ書けるかもしれない
class TreeVisitor:
    def __init__(self, command):
        self.command = command
    def visit(self, node):
        if node is None:
            return
        self.command.whenPreVisit(node)
        if node.left is not None:
            node.left.accept(self)
        self.command.whenInVisit(node)
        if node.right is not None:
            node.right.accept(self)
        self.command.whenPostVisit(node)

class Command:
    def __init__(self):
        pass
    def whenPreVisit(self, node):
        pass
    def whenInVisit(self, node):
        pass
    def whenPostVisit(self, node):
        pass
    def getResult(self):
        return None

class PreOrderKeyCollector(Command):
    def __init__(self):
        super().__init__()
        self.keys = []
    def whenPreVisit(self, node):
        self.keys.append(node.key)
    def getResult(self):
        return ' ' + ' '.join([str(x) for x in self.keys])

class InOrderKeyCollector(Command):
    def __init__(self):
        super().__init__()
        self.keys = []
    def whenInVisit(self, node):
        self.keys.append(node.key)
    def getResult(self):
        return ' ' + ' '.join([str(x) for x in self.keys])

class MaxLevelCounter(Command):
    def __init__(self):
        super().__init__()
        self.level = 0
        self.max_level = 0
    def whenPreVisit(self, node):
        self.level += 1
    def whenInVisit(self, node):
        if self.level > self.max_level:
            self.max_level = self.level
    def whenPostVisit(self, node):
        self.level -= 1
    def getResult(self):
        return self.max_level

class NodeCounter(Command):
    def __init__(self):
        super().__init__()
        self.count = 0
    def whenInVisit(self, node):
        self.count += 1
    def getResult(self):
        return self.count