flat-tree/flat_tree/accessor.py

"""An accessor for navigating flat trees."""

__all__ = ["FlatTreeAccessor"]

from typing import List, Optional


class FlatTreeAccessor:
    """A flat tree accessor."""

    def index(self, depth: int, offset: int) -> int:
        """The tree index specified by the depth and offset.

        :param depth: The depth of the tree
        :param offset: The offset from left hand side of the tree
        """
        return ((1 + (2 * offset)) * (2 ** depth)) - 1

    def offset(self, index: int, depth: Optional[int] = None) -> int:
        """The offset of given index from the left hand side of the tree.

        :param index: The tree index
        :param depth: The depth of the tree
        """
        if not (index & 1):
            return int(index / 2)

        if depth is None:
            depth = self.depth(index)

        return int((((index + 1) / (2 ** depth)) - 1) / 2)

    def depth(self, index: int) -> int:
        """The depth of the given index in the tree.

        :param index: The tree index
        """
        depth = 0

        index += 1
        while not (index & 1):
            depth += 1
            index = index >> 1

        return depth

    def parent(self, index: int, depth: Optional[int] = None) -> int:
        """The index of the parent relative to the given index.

        :param index: The index relative to the parent
        :param depth: The depth of the index
        """
        if depth is None:
            depth = self.depth(index)

        offset = self.offset(index, depth)

        return self.index(depth + 1, int((offset - (offset & 1)) / 2))

    def sibling(self, index: int, depth: Optional[int] = None) -> int:
        """The index of the sibling relative to the given index.

        :param index: The index relative to the sibling
        :param depth: The depth of the index
        """
        if depth is None:
            depth = self.depth(index)

        offset = self.offset(index, depth)
        offset = offset - 1 if (offset & 1) else offset + 1

        return self.index(depth, offset)

    def children(self, index: int, depth: Optional[int] = None) -> List[int]:
        """All children relative to the given index.

        :param index: The parent index
        :param depth: The depth of the index
        """
        if not (index & 1):
            return []

        if not depth:
            depth = self.depth(index)

        offset = self.offset(index, depth) * 2

        return [
            self.index((depth - 1), offset),
            self.index((depth - 1), (offset + 1)),
        ]

    def spans(self, index: int, depth: Optional[int] = None) -> List[int]:
        """The span of the tree.

        :param index: The index of the root
        :param depth: The depth of the index
        """
        if not (index & 1):
            return [index, index]

        if not depth:
            depth = self.depth(index)

        offset = self.offset(index, depth)
        width = 2 ** (depth + 1)

        return [(offset * width), ((offset + 1) * width) - 2]

    def left_span(self, index: int, depth: Optional[int] = None) -> int:
        """The leftmost span of the tree.

        :param index: The index of the tree root
        :param depth: The depth of the index
        """
        if not (index & 1):
            return index

        if not depth:
            depth = self.depth(index)

        return self.offset(index, depth) * (2 ** (depth + 1))

    def right_span(self, index: int, depth: Optional[int] = None) -> int:
        """The rightmost span of the tree.

        :param index: The index of the tree root
        :param depth: The depth of the index
        """
        if not (index & 1):
            return index

        if not depth:
            depth = self.depth(index)

        return (self.offset(index, depth) + 1) * (2 ** (depth + 1)) - 2

    def count(self, index: int, depth: Optional[int] = None) -> int:
        """The number of nodes a tree contains.

        :param index: The index of the root of the tree
        :param depth: The depth of the root of the tree
        """
        if not (index & 1):
            return 1

        if not depth:
            depth = self.depth(index)

        return (2 ** (depth + 1)) - 1

    def full_roots(self, index: int) -> List[int]:
        """All full roots within the tree.

        :param index: The index of the root of the tree
        """
        if index & 1:
            message = "Roots only available for tree depth 0"
            raise ValueError(message)

        roots: List[int] = []

        index = int(index / 2)
        offset, factor = 0, 1

        while True:
            if not index:
                return roots

            while (factor * 2) <= index:
                factor = factor * 2

            roots.append((offset + factor) - 1)

            offset = (offset + 2) * factor
            index = index - factor
            factor = 1

    def left_child(self, index: int, depth: Optional[int] = None) -> int:
        """The left child of the given index.

        :param index: The index of the tree
        :param depth: The depth of the tree
        """
        if not (index & 1):
            return -1

        if not depth:
            depth = self.depth(index)

        return self.index((depth - 1), (self.offset(index, depth) * 2))

    def right_child(self, index: int, depth: Optional[int] = None) -> int:
        """The right child of the given index.

        :param index: The index of the tree
        :param depth: The depth of the tree
        """
        if not (index & 1):
            return -1

        if not depth:
            depth = self.depth(index)

        return self.index((depth - 1), (1 + (self.offset(index, depth) * 2)))