rooted_unbounded_tree.py

"""
Rooted tree with unbounded branching, left child - right sibling representation backed by 4 arrays.
"""


class UnboundedTree(object):

    def __init__(self):
        self.root = 0
        self.storage = [
            [None],     # parent
            [],         # keys
            [None],     # left child
            [None]      # right sibling
        ]

    def add(self, key, parent=0):
        self.storage[0].append(parent)
        self.storage[1].append(key)
        self.storage[2].append(None)
        self.storage[3].append(None)
        element_index = len(self.storage[0]) - 1
        assert parent <= len(self.storage[2]) - 1
        left_child = self.storage[2][parent]
        if left_child is None:
            # adding as a left child
            self.storage[2][parent] = element_index
        else:
            # left child exists, adding right sibling
            self.storage[3][left_child] = element_index

    def traverse_recursive(self):
        """Returns all keys stored in the current tree. 10.4-4"""
        keys = []

        def visit_children(node):
            left_child = self.storage[2][node]
            if left_child is None:
                return
            keys.append(self.storage[1][left_child])
            visit_children(self.storage[2][left_child])
            right_sibling = self.storage[3][left_child]
            while right_sibling is not None:
                assert self.storage[0][right_sibling] == node
                keys.append(self.storage[1][right_sibling])
                visit_children(right_sibling)
                right_sibling = self.storage[3][right_sibling]
        visit_children(self.root)
        return keys

    def has_key(self, key):
        return self._find(key) is not None

    def _find(self, key):
        """10.4-3"""
        stack = [self.root]
        while stack:
            node = stack.pop()
            if self.storage[1][node] == key:
                return node
            left_child = self.storage[2][node]
            if left_child is None:
                continue
            stack.append(left_child)
            right_sibling = self.storage[3][left_child]
            while right_sibling is not None:
                assert self.storage[0][right_sibling] == node
                stack.append(right_sibling)
                right_sibling = self.storage[3][right_sibling]
        return None