btree.py

#!/usr/bin/env python3

""" Binary search tree with method to invert in-place. """

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

class Btree(object):
    """ Binary search tree."""

    def __init__(self):
        self.root = None

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

        self.root = self._insert(value, self.root)
        return

    def _insert(self, value, node):
        if node is None:
            node = Node(value)
        elif value < node.value:
            node.left = self._insert(value, node.left)
        elif value > node.value:
            node.right = self._insert(value, node.right)

        return node

    def walk(self):
        if self.root is None:
            out = []
        else:
            out = self._walk(self.root, [])
        return out

    def _walk(self, node, acc):
        """ In-order traversal."""

        if node is None:
            return acc

        acc = self._walk(node.left, acc)
        acc = acc + [node.value]
        acc = self._walk(node.right, acc)
        return acc

    def __contains__(self, x):
        return self._search(self.root, x)

    def _search(self, node, x):
        if node is None:
            return False

        if target == node.value:
            return True
        elif target < node.value:
            return self._search(node.left, x)
        else:
            return self._search(node.right, x)

class InvertableBtree(Btree):
    def invert(self):
        """ Invert the tree so that traversing it in-order gives the
        elements in descending (reverse sorted) order. Updates the
        tree in-place.
        """

        self.root = self._invert(self.root)
        return

    def _invert(self, node):
        if node is None:
            return node

        node.left, node.right = node.right, node.left

        node.left = self._invert(node.left)
        node.right = self._invert(node.right)
        return node


if __name__ == '__main__':
    xs = [10, 2, 15, 1, 5, 11, 20]

    tree = InvertableBtree()

    for x in xs:
        tree.insert(x)

    assert tree.walk() == sorted(xs)

    tree.invert()
    assert tree.walk() == sorted(xs, reverse=True)