Binary search trees

Base class and simple binary search tree

BaseBinarySearchTree

Base class for binary search trees.

Source code in tralda/datastructures/bst/base.py

class BaseBinarySearchTree:
    """Base class for binary search trees."""

    __slots__ = ("root", "_temp_attributes")

    node_class = BinaryNode
    iterator_class: Iterator[Any] = BinaryTreeIterator

    def __init__(self) -> None:
        """Initialize the balanced binary search tree."""
        self.root: BinaryNode | None = None

        # for temporarily storing the node attributes for insertion/popping in recursive functions
        self._temp_attributes: Any | None = None

    def __iter__(self) -> Iterator[Any]:
        """An iterator for the binary search tree.

        Returns:
            An iterator for the binary search tree.
        """
        return self.iterator_class(self)

    def __next__(self) -> None:
        raise TypeError(
            f"'{type(self).__name__}' object is not an iterator; "
            "use iter() to obtain a tree iterator"
        )

    def __bool__(self) -> bool:
        """Return whether the tree is non-empty.

        Returns:
            Whether the tree contains elements or not.
        """
        return bool(self.root)

    def __len__(self) -> int:
        """Number of elements in the tree.

        Returns:
            Number of elements.
        """
        return self.root.size if self.root else 0

    def __contains__(self, item: Any) -> bool:
        """Return whether the tree contains the item.

        Args:
            item: The item for which to check whether it is contained in the tree.

        Returns:
            Whether the tree contains the item.
        """
        return self._find(item) is not None

    def __getitem__(self, idx: int) -> Any:
        """Return the element at the index.

        Same as 'key_at_index(index)'.

        Args:
            idx: The index of the key to get.

        Returns:
            The key of the node at the index.
        """
        return self._node_at_index(idx).key

    def key_at_index(self, idx: int) -> Any:
        """Return the key at the index.

        Args:
            idx: The index of the key to get.

        Returns:
            The key of the node at the index.
        """
        return self._node_at_index(idx).key

    def _validate_index(self, idx: int) -> int:
        """Return the node at the index.

        Args:
            idx: The index to be validated.

        Returns:
            The input index or the corresponding positive index if the input was negative.

        Raises:
            IndexError: If the index is out of bounds.
        """
        if not self.root:
            raise IndexError(f"index {idx} is out of range")

        if idx < 0:
            if idx < -self.root.size:
                raise IndexError(f"index {idx} is out of range")
            else:
                idx += self.root.size

        if idx >= self.root.size:
            raise IndexError(f"index {idx} is out of range")

        return idx

    def _node_at_index(self, idx: int) -> BinaryNode:
        """Return the node at the index.

        Args:
            idx: The index.

        Returns:
            The node instance at the index.

        Raises:
            IndexError: If the index is out of bounds.
            RuntimeError: If the index seems valid but the node could not be found. A corrupted
                integrity of the tree datastructure could be the reason.
        """
        idx = self._validate_index(idx)
        current = self.root
        current_sum = 0

        while current:
            current_idx = current_sum + current.left_size()
            if idx == current_idx:
                return current
            elif idx < current_idx:
                current = current.left
            else:
                current = current.right
                current_sum = current_idx + 1

        raise RuntimeError(f"could not find node with index {idx}")

    def add(self, item: Any) -> None:
        """Add an item if not yet present.

        Args:
            item: The new item to be inserted.
        """
        try:
            self.insert(item)
        except (KeyError, ValueError):
            pass

    def insert(self, key: Any) -> None:
        """Insert an item.

        This function will throw a ValueError if the key is already present. If you do not want
        this behavior, use the function add() instead.

        Args:
            key: The new item to be inserted.

        Raises:
            KeyError: If the key already exists.
        """
        self._temp_attributes = (key,)
        self._insert_key(key)
        self._temp_attributes = None

    def remove(self, key: Any) -> None:
        """Remove a key from the tree.

        Args:
            key: The new item to be removed.

        Raises:
            KeyError: If the key is not in the tree.
        """
        self._delete_key(key)

    def discard(self, key: Any) -> None:
        """Remove a key from the tree if present.

        Args:
            key: The new item to be removed.
        """
        try:
            self._delete_key(key)
        except KeyError:
            pass

    def pop(self) -> Any:
        """Remove and return the greatest item.

        Returns:
            The greatest item in the tree.

        Raises:
            IndexError: If the tree is empty.
        """
        if not self.root:
            raise IndexError("pop from empty tree")

        return self.pop_at_index(len(self) - 1)

    def clear(self) -> None:
        """Removes all items from the tree."""
        self.root = None

    def difference_update(self, items: Iterable[Any]) -> None:
        """Discard all elements in the collection.

        Args:
            items: The items to be discarded.
        """
        for item in items:
            self.discard(item)

    def remove_at_index(self, idx: int) -> None:
        """Remove node at the index.

        Args:
            idx: The index of the element to be removed.

        Raises:
            IndexError: If the index is out of bounds.
        """
        self.pop_at_index(idx)

    def pop_at_index(self, idx: int) -> Any:
        """Remove item at the index.

        Args:
            idx: The index of the element to be removed and returned.

        Returns:
            The item at the index.

        Raises:
            IndexError: If the index is out of bounds.
        """
        idx = self._validate_index(idx)
        self._pop_at_index(idx)

        to_pop = tuple(self._temp_attributes)
        self._temp_attributes = None

        if len(to_pop) == 1:
            to_pop = to_pop[0]

        return to_pop

    def _find(self, key: Any) -> BinaryNode | None:
        """Find the node for the specified key.

        Args:
            key: The key to be searched.

        Returns:
            The corresponding tree node or None if the key was not found.
        """
        if not self.root:
            return None

        current = self.root
        while current:
            if key == current.key:
                return current
            elif key < current.key:
                current = current.left
            else:
                current = current.right

    def _insert_key(self, key: Any) -> None:
        """Insert a key into the tree if not already present.

        Args:
            key: The key to be inserted.

        Raises:
            KeyError: If the key already exists.
            NotImplementedError: If the child class does not implement this method.
        """
        raise NotImplementedError("not implemented for base class")

    def _delete_key(self, key: Any) -> None:
        """Delete a key.

        Args:
            key: The key to be deleted.

        Raises:
            KeyError: If the key does not exist.
            NotImplementedError: If the child class does not implement this method.
        """
        raise NotImplementedError("not implemented for base class")

    def _pop_at_index(self, idx: int) -> None:
        """Remove item at the index and return it.

        Args:
            idx: The index of the element to be removed.

        Raises:
            IndexError: If the index is out of bounds.
            NotImplementedError: If the child class does not implement this method.
        """
        raise NotImplementedError("not implemented for base class")

    def _inorder_traversal(self) -> Iterator[BinaryNode]:
        """Generator for the nodes in a in-order traversal of the tree.

        Yields:
            All nodes of the tree in in-order.
        """

        def _inorder(node: BinaryNode) -> Iterator[BinaryNode]:
            if node.left:
                yield from _inorder(node.left)
            yield node
            if node.right:
                yield from _inorder(node.right)

        if self.root:
            yield from _inorder(self.root)
        else:
            yield from []

    def _postorder_traversal(self) -> Iterator[BinaryNode]:
        """Generator for the nodes in a post-order traversal of the tree.

        Yields:
            All nodes of the tree in post-order.
        """

        def _postorder(node: BinaryNode) -> Iterator[BinaryNode]:
            if node.left:
                yield from _postorder(node.left)
            if node.right:
                yield from _postorder(node.right)
            yield node

        if self.root:
            yield from _postorder(self.root)
        else:
            yield from []

    def copy(self) -> BaseBinarySearchTree:
        """Copy the tree.

        Returns:
            A copy of the tree.
        """
        tree_copy = self.__class__()
        if self.root:
            tree_copy.root = self._copy_subtree(self.root)

        return tree_copy

    def _copy_subtree(self, node: BinaryNode) -> BinaryNode:
        """Copy a subtree.

        Args:
            node: The node whose subtree to copy.

        Returns:
            The copied node (= the root of the copied subtree).
        """
        node_copy = node.copy()
        if node.left:
            node_copy.left = self._copy_subtree(node.left)
        if node.right:
            node_copy.right = self._copy_subtree(node.right)

        return node_copy

    def to_newick(self) -> str:
        """Newick representation of the tree.

        Intended for testing purpose.

        Returns:
            A Newick representation of the tree.
        """

        def _newick(node):
            if not (node.left or node.right):
                return str(node.key)
            else:
                if node.left and node.right:
                    s = f"({_newick(node.left)},{_newick(node.right)})"
                elif node.left:
                    s = f"({_newick(node.left)},-)"
                elif node.right:
                    s = f"(-,{_newick(node.right)})"
                else:
                    s = ""
                return s + str(node.key)

        return _newick(self.root) if self.root else ""

    def check_integrity(self, verbose: bool = False) -> bool:
        """Integrity check of the tree.

        Checks whether the size and height is correct in all subtrees. Intended for debugging and
        testing purpose.

        Args:
            verbose: If True print where the integrity check has failed.

        Returns:
            Whether all integrity checks have been passed.
        """
        for node in self._postorder_traversal():
            height_left, height_right, size = 0, 0, 1

            if node.left:
                height_left = node.left.height
                size += node.left.size

            if node.right:
                height_right = node.right.height
                size += node.right.size

            height = 1 + max(height_left, height_right)
            if node.height != height:
                if verbose:
                    print(
                        f"height of node {node} is incorrect (expected {height}, got {node.height})"
                    )
                return False

            if node.size != size:
                if verbose:
                    print(f"size of node {node} is incorrect (expected {size}, got {node.size})")
                return False

        return True

__init__() -> None

Initialize the balanced binary search tree.

Source code in tralda/datastructures/bst/base.py

def __init__(self) -> None:
    """Initialize the balanced binary search tree."""
    self.root: BinaryNode | None = None

    # for temporarily storing the node attributes for insertion/popping in recursive functions
    self._temp_attributes: Any | None = None

__iter__() -> Iterator[Any]

An iterator for the binary search tree.

Returns:

Type	Description
Iterator[Any]	An iterator for the binary search tree.

Source code in tralda/datastructures/bst/base.py

def __iter__(self) -> Iterator[Any]:
    """An iterator for the binary search tree.

    Returns:
        An iterator for the binary search tree.
    """
    return self.iterator_class(self)

__bool__() -> bool

Return whether the tree is non-empty.

Returns:

Type	Description
bool	Whether the tree contains elements or not.

Source code in tralda/datastructures/bst/base.py

def __bool__(self) -> bool:
    """Return whether the tree is non-empty.

    Returns:
        Whether the tree contains elements or not.
    """
    return bool(self.root)

__len__() -> int

Number of elements in the tree.

Returns:

Type	Description
int	Number of elements.

Source code in tralda/datastructures/bst/base.py

def __len__(self) -> int:
    """Number of elements in the tree.

    Returns:
        Number of elements.
    """
    return self.root.size if self.root else 0

__contains__(item: Any) -> bool

Return whether the tree contains the item.

Parameters:

Name	Type	Description	Default
item	Any	The item for which to check whether it is contained in the tree.	required

Returns:

Type	Description
bool	Whether the tree contains the item.

Source code in tralda/datastructures/bst/base.py

def __contains__(self, item: Any) -> bool:
    """Return whether the tree contains the item.

    Args:
        item: The item for which to check whether it is contained in the tree.

    Returns:
        Whether the tree contains the item.
    """
    return self._find(item) is not None

__getitem__(idx: int) -> Any

Return the element at the index.

Same as 'key_at_index(index)'.

Parameters:

Name	Type	Description	Default
idx	int	The index of the key to get.	required

Returns:

Type	Description
Any	The key of the node at the index.

Source code in tralda/datastructures/bst/base.py

def __getitem__(self, idx: int) -> Any:
    """Return the element at the index.

    Same as 'key_at_index(index)'.

    Args:
        idx: The index of the key to get.

    Returns:
        The key of the node at the index.
    """
    return self._node_at_index(idx).key

key_at_index(idx: int) -> Any

Return the key at the index.

Parameters:

Name	Type	Description	Default
idx	int	The index of the key to get.	required

Returns:

Type	Description
Any	The key of the node at the index.

Source code in tralda/datastructures/bst/base.py

def key_at_index(self, idx: int) -> Any:
    """Return the key at the index.

    Args:
        idx: The index of the key to get.

    Returns:
        The key of the node at the index.
    """
    return self._node_at_index(idx).key

add(item: Any) -> None

Add an item if not yet present.

Parameters:

Name	Type	Description	Default
item	Any	The new item to be inserted.	required

Source code in tralda/datastructures/bst/base.py

def add(self, item: Any) -> None:
    """Add an item if not yet present.

    Args:
        item: The new item to be inserted.
    """
    try:
        self.insert(item)
    except (KeyError, ValueError):
        pass

insert(key: Any) -> None

Insert an item.

This function will throw a ValueError if the key is already present. If you do not want this behavior, use the function add() instead.

Parameters:

Name	Type	Description	Default
key	Any	The new item to be inserted.	required

Raises:

Type	Description
KeyError	If the key already exists.

Source code in tralda/datastructures/bst/base.py

def insert(self, key: Any) -> None:
    """Insert an item.

    This function will throw a ValueError if the key is already present. If you do not want
    this behavior, use the function add() instead.

    Args:
        key: The new item to be inserted.

    Raises:
        KeyError: If the key already exists.
    """
    self._temp_attributes = (key,)
    self._insert_key(key)
    self._temp_attributes = None

remove(key: Any) -> None

Remove a key from the tree.

Parameters:

Name	Type	Description	Default
key	Any	The new item to be removed.	required

Raises:

Type	Description
KeyError	If the key is not in the tree.

Source code in tralda/datastructures/bst/base.py

def remove(self, key: Any) -> None:
    """Remove a key from the tree.

    Args:
        key: The new item to be removed.

    Raises:
        KeyError: If the key is not in the tree.
    """
    self._delete_key(key)

discard(key: Any) -> None

Remove a key from the tree if present.

Parameters:

Name	Type	Description	Default
key	Any	The new item to be removed.	required

Source code in tralda/datastructures/bst/base.py

def discard(self, key: Any) -> None:
    """Remove a key from the tree if present.

    Args:
        key: The new item to be removed.
    """
    try:
        self._delete_key(key)
    except KeyError:
        pass

pop() -> Any

Remove and return the greatest item.

Returns:

Type	Description
Any	The greatest item in the tree.

Raises:

Type	Description
IndexError	If the tree is empty.

Source code in tralda/datastructures/bst/base.py

def pop(self) -> Any:
    """Remove and return the greatest item.

    Returns:
        The greatest item in the tree.

    Raises:
        IndexError: If the tree is empty.
    """
    if not self.root:
        raise IndexError("pop from empty tree")

    return self.pop_at_index(len(self) - 1)

clear() -> None

Removes all items from the tree.

Source code in tralda/datastructures/bst/base.py

def clear(self) -> None:
    """Removes all items from the tree."""
    self.root = None

difference_update(items: Iterable[Any]) -> None

Discard all elements in the collection.

Parameters:

Name	Type	Description	Default
items	Iterable[Any]	The items to be discarded.	required

Source code in tralda/datastructures/bst/base.py

def difference_update(self, items: Iterable[Any]) -> None:
    """Discard all elements in the collection.

    Args:
        items: The items to be discarded.
    """
    for item in items:
        self.discard(item)

remove_at_index(idx: int) -> None

Remove node at the index.

Parameters:

Name	Type	Description	Default
idx	int	The index of the element to be removed.	required

Raises:

Type	Description
IndexError	If the index is out of bounds.

Source code in tralda/datastructures/bst/base.py

def remove_at_index(self, idx: int) -> None:
    """Remove node at the index.

    Args:
        idx: The index of the element to be removed.

    Raises:
        IndexError: If the index is out of bounds.
    """
    self.pop_at_index(idx)

pop_at_index(idx: int) -> Any

Remove item at the index.

Parameters:

Name	Type	Description	Default
idx	int	The index of the element to be removed and returned.	required

Returns:

Type	Description
Any	The item at the index.

Raises:

Type	Description
IndexError	If the index is out of bounds.

Source code in tralda/datastructures/bst/base.py

def pop_at_index(self, idx: int) -> Any:
    """Remove item at the index.

    Args:
        idx: The index of the element to be removed and returned.

    Returns:
        The item at the index.

    Raises:
        IndexError: If the index is out of bounds.
    """
    idx = self._validate_index(idx)
    self._pop_at_index(idx)

    to_pop = tuple(self._temp_attributes)
    self._temp_attributes = None

    if len(to_pop) == 1:
        to_pop = to_pop[0]

    return to_pop

copy() -> BaseBinarySearchTree

Copy the tree.

Returns:

Type	Description
BaseBinarySearchTree	A copy of the tree.

Source code in tralda/datastructures/bst/base.py

def copy(self) -> BaseBinarySearchTree:
    """Copy the tree.

    Returns:
        A copy of the tree.
    """
    tree_copy = self.__class__()
    if self.root:
        tree_copy.root = self._copy_subtree(self.root)

    return tree_copy

to_newick() -> str

Newick representation of the tree.

Intended for testing purpose.

Returns:

Type	Description
str	A Newick representation of the tree.

Source code in tralda/datastructures/bst/base.py

def to_newick(self) -> str:
    """Newick representation of the tree.

    Intended for testing purpose.

    Returns:
        A Newick representation of the tree.
    """

    def _newick(node):
        if not (node.left or node.right):
            return str(node.key)
        else:
            if node.left and node.right:
                s = f"({_newick(node.left)},{_newick(node.right)})"
            elif node.left:
                s = f"({_newick(node.left)},-)"
            elif node.right:
                s = f"(-,{_newick(node.right)})"
            else:
                s = ""
            return s + str(node.key)

    return _newick(self.root) if self.root else ""

check_integrity(verbose: bool = False) -> bool

Integrity check of the tree.

Checks whether the size and height is correct in all subtrees. Intended for debugging and testing purpose.

Parameters:

Name	Type	Description	Default
verbose	bool	If True print where the integrity check has failed.	False

Returns:

Type	Description
bool	Whether all integrity checks have been passed.

Source code in tralda/datastructures/bst/base.py

def check_integrity(self, verbose: bool = False) -> bool:
    """Integrity check of the tree.

    Checks whether the size and height is correct in all subtrees. Intended for debugging and
    testing purpose.

    Args:
        verbose: If True print where the integrity check has failed.

    Returns:
        Whether all integrity checks have been passed.
    """
    for node in self._postorder_traversal():
        height_left, height_right, size = 0, 0, 1

        if node.left:
            height_left = node.left.height
            size += node.left.size

        if node.right:
            height_right = node.right.height
            size += node.right.size

        height = 1 + max(height_left, height_right)
        if node.height != height:
            if verbose:
                print(
                    f"height of node {node} is incorrect (expected {height}, got {node.height})"
                )
            return False

        if node.size != size:
            if verbose:
                print(f"size of node {node} is incorrect (expected {size}, got {node.size})")
            return False

    return True

BinarySearchTree

Bases: BaseBinarySearchTree

Simple binary search tree.

Source code in tralda/datastructures/bst/simple.py

class BinarySearchTree(BaseBinarySearchTree):
    """Simple binary search tree."""

    __slots__ = ()

    node_class = BinaryNode
    iterator_class: Iterator[Any] = BinaryTreeIterator

    def _insert_key(self, key: Any) -> None:
        """Insert a key into the tree if not already present.

        Args:
            key: The key to be inserted.
        """
        self.root = self._recursive_insert(key, self.root)

    def _recursive_insert(self, key: Any, node: BinaryNode | None) -> BinaryNode:
        """Recursive insertion and rebalancing.

        Args:
            key: The key to be inserted.
            node: The node below which to insert the new key or None if the node should be inserted
                at this position.

        Returns:
            The provided node or the newly created node if None was provided.

        Raises:
            KeyError: If the key already exists.
        """
        if node is None:
            return self.node_class(key)
        elif key < node.key:
            node.left = self._recursive_insert(key, node.left)
        elif key > node.key:
            node.right = self._recursive_insert(key, node.right)
        else:
            raise KeyError(f"key {key} already exists")

        node.update()

        return node

    def _delete_key(self, key: Any) -> None:
        """Delete a key if present.

        Args:
            key: The key to be deleted.
        """
        self.root = self._recursive_delete(key, self.root)

    def _recursive_delete(self, key: Any, node: BinaryNode | None) -> BinaryNode | None:
        """Recursive deletion.

        Args:
            key: The key to be deleted.
            node: The node below which the key to be deleted is located.

        Returns:
            The provided node, a node that is moved up, or None if a leaf node was deleted.

        Raises:
            KeyError: If the key is not in the tree.
        """
        if node is None:
            raise KeyError(f"key {key} not found")

        if key < node.key:
            node.left = self._recursive_delete(key, node.left)
        elif key > node.key:
            node.right = self._recursive_delete(key, node.right)
        else:
            if node.left is None:
                node = node.right
            elif node.right is None:
                node = node.left
            else:
                subst_node = node.right.smallest_in_subtree()
                subst_node.copy_attributes_to_node(node)

                # now find and delete subst_node
                node.right = self._recursive_delete(node.key, node.right)

        if node is not None:
            node.update()

        return node

    def _pop_at_index(self, idx: int) -> None:
        """Remove item at the index and return it.

        Args:
            idx: The index of the element to be removed.
        """
        self.root = self._recursive_pop_at_index(idx, self.root)

    def _recursive_pop_at_index(self, idx: int, node: BinaryNode) -> BinaryNode | None:
        """Recursive deletion of a node at the index.

        Args:
            idx: The index of the element to be popped.
            node: The node below which the node with the specified index is located (or equal to
                this node).

        Returns:
            The provided node, a node that is moved up, or None if a leaf node was deleted.

        Raises:
            IndexError: If the key is not in the tree.
        """
        if node is None:
            raise IndexError(f"could not find node with index {idx}")

        node_idx = node.left_size()

        if idx == node_idx:
            self._temp_attributes = node.get_attributes()
            return self._recursive_delete(node.key, node)
        elif idx < node_idx:
            node.left = self._recursive_pop_at_index(idx, node.left)
        else:
            node.right = self._recursive_pop_at_index(idx - node_idx - 1, node.right)

        if node is not None:
            node.update()

        return node

---

AVL-tree-based data structures

TreeSet

Bases: BaseBinarySearchTree

AVL tree implementation of a sorted set.

Source code in tralda/datastructures/bst/avl.py

class TreeSet(BaseBinarySearchTree):
    """AVL tree implementation of a sorted set."""

    __slots__ = ()

    node_class = BinaryNode
    iterator_class: Iterator[Any] = BinaryTreeIterator

    def check_integrity(self, verbose: bool = False) -> bool:
        """Integrity check of the tree.

        Checks whether all children have a correct parent reference and the size and height is
        correct in all subtrees. Additionally, the AVL property is checked. Intended for debugging
        and testing purpose.

        Args:
            verbose: If True print where the integrity check has failed.

        Returns:
            Whether all integrity checks have been passed.
        """
        if not super().check_integrity(verbose=verbose):
            return False

        for node in self._inorder_traversal():
            if abs(node.balance()) > 1:
                if verbose:
                    print(f"node {node} is unbalanced")
                return False

        return True

    def _rotate_right(self, node: BinaryNode) -> BinaryNode:
        """Perform a right rotation on the node.

        Args:
            node: The node on which to perform a right rotation.

        Returns:
            The former left child of the node which is its parent after the rotation.
        """
        left_child = node.left
        node.left = left_child.right
        left_child.right = node

        node.update()
        left_child.update()

        return left_child

    def _rotate_left(self, node: BinaryNode) -> BinaryNode:
        """Perform a left rotation on the node.

        Args:
            node: The node on which to perform a left rotation.

        Returns:
            The former right child of the node which is its parent after the rotation.
        """
        right_child = node.right
        node.right = right_child.left
        right_child.left = node

        node.update()
        right_child.update()

        return right_child

    def _rebalance(self, node: BinaryNode) -> BinaryNode:
        """Rebalance the node.

        Args:
            node: The node to be rebalanced.

        Returns:
            The root of the tree after rebalancing.
        """
        node.update()
        balance = node.balance()

        if balance > 1:
            if node.left.balance() >= 0:
                # single rotation needed
                node = self._rotate_right(node)
            else:
                # double rotation needed
                node.left = self._rotate_left(node.left)
                node = self._rotate_right(node)
        elif balance < -1:
            if node.right.balance() <= 0:
                # single rotation needed
                node = self._rotate_left(node)
            else:
                # double rotation needed
                node.right = self._rotate_right(node.right)
                node = self._rotate_left(node)

        return node

    def _insert_key(self, key: Any) -> None:
        """Insert a key into the tree if not already present.

        Args:
            key: The key to be inserted.
        """
        self.root = self._insert_and_rebalance(key, self.root)

    def _insert_and_rebalance(self, key: Any, node: BinaryNode | None) -> BinaryNode:
        """Recursive insertion and rebalancing.

        Args:
            key: The key to be inserted.
            node: The node under which the key must be inserted.

        Returns:
            The root of the subtree after inserting the key and rebalancing.

        Raises:
            KeyError: If the key already exists.
        """
        if node is None:
            return self.node_class(*self._temp_attributes)

        if key < node.key:
            node.left = self._insert_and_rebalance(key, node.left)
        elif key > node.key:
            node.right = self._insert_and_rebalance(key, node.right)
        else:
            raise KeyError(f"key {key} already exists")

        return self._rebalance(node)

    def _delete_key(self, key: Any) -> None:
        """Delete a key if present.

        Args:
            key: The key to be deleted.
        """
        self.root = self._delete_and_rebalance(key, self.root)

    def _delete_and_rebalance(self, key: Any, node: BinaryNode | None) -> BinaryNode:
        """Recursive deletion and rebalancing.

        Args:
            key: The key to be deleted.
            node: The node below which the key to be deleted is located.

        Returns:
            The root of the subtree after deletion and rebalancing.

        Raises:
            KeyError: If the key is not in the tree.
        """
        if node is None:
            raise KeyError(f"key {key} not found")

        if key < node.key:
            node.left = self._delete_and_rebalance(key, node.left)
        elif key > node.key:
            node.right = self._delete_and_rebalance(key, node.right)
        else:
            if node.left is None:
                node = node.right
            elif node.right is None:
                node = node.left
            else:
                subst_node = node.right.smallest_in_subtree()
                subst_node.copy_attributes_to_node(node)
                # now find and delete subst_node
                node.right = self._delete_and_rebalance(node.key, node.right)

        if node is not None:
            node = self._rebalance(node)

        return node

    def _pop_at_index(self, idx: int) -> None:
        """Remove item at the index.

        Args:
            idx: The index of the element to be removed.
        """
        self.root = self._pop_at_index_and_rebalance(idx, self.root)

    def _pop_at_index_and_rebalance(self, idx: int, node: BinaryNode) -> BinaryNode:
        """Recursive deletion of a node at the index."""
        if node is None:
            raise ValueError(f"could not find node with index {idx}")

        node_idx = node.left_size()

        if idx == node_idx:
            self._temp_attributes = node.get_attributes()
            return self._delete_and_rebalance(node.key, node)
        elif idx < node_idx:
            node.left = self._pop_at_index_and_rebalance(idx, node.left)
        else:
            node.right = self._pop_at_index_and_rebalance(idx - node_idx - 1, node.right)

        if node is not None:
            node = self._rebalance(node)

        return node

check_integrity(verbose: bool = False) -> bool

Integrity check of the tree.

Checks whether all children have a correct parent reference and the size and height is correct in all subtrees. Additionally, the AVL property is checked. Intended for debugging and testing purpose.

Parameters:

Name	Type	Description	Default
verbose	bool	If True print where the integrity check has failed.	False

Returns:

Type	Description
bool	Whether all integrity checks have been passed.

Source code in tralda/datastructures/bst/avl.py

def check_integrity(self, verbose: bool = False) -> bool:
    """Integrity check of the tree.

    Checks whether all children have a correct parent reference and the size and height is
    correct in all subtrees. Additionally, the AVL property is checked. Intended for debugging
    and testing purpose.

    Args:
        verbose: If True print where the integrity check has failed.

    Returns:
        Whether all integrity checks have been passed.
    """
    if not super().check_integrity(verbose=verbose):
        return False

    for node in self._inorder_traversal():
        if abs(node.balance()) > 1:
            if verbose:
                print(f"node {node} is unbalanced")
            return False

    return True

TreeDict

Bases: TreeSet

Source code in tralda/datastructures/bst/avl.py

class TreeDict(TreeSet):
    __slots__ = ()

    node_class = TreeDictNode
    iterator_class: Iterator[Any] = TreeDictIterator

    def __getitem__(self, key: Any) -> Any:
        """Return the value associated with the key.

        Overrides the method of the base class where the key at a specified index is returned.

        Args:
            key: The key.

        Returns:
            The value associated with the key.

        Raises:
            KeyError: If the key does not exist.
        """
        node = self._find(key)

        if not node:
            raise KeyError(f"key {key} is not in the TreeDict")

        return node.value

    def get(self, key: Any, default: Any = None) -> Any:
        """Return the value associated with the key or a default value if the key does not exist.

        Args:
            key: The key for which to get the associated value.
            default: The default to be returned if the key does not exist.

        Returns:
            The value associated with the key or the default value.
        """
        node = self._find(key)

        if node:
            return node.value
        else:
            return default

    def keys(self) -> Iterator[Any]:
        """Iterator for the keys.

        Returns:
            An iterator for the keys.
        """
        return self.iterator_class(self, mode=TreeDictIteratorMode.KEY)

    def values(self) -> Iterator[Any]:
        """Iterator for the values.

        Returns:
            An iterator for the values.
        """
        return self.iterator_class(self, mode=TreeDictIteratorMode.VALUE)

    def items(self) -> Iterator[tuple[Any, Any]]:
        """Iterator for the key-value pairs.

        Returns:
            An iterator for the key-value pairs.
        """
        return self.iterator_class(self, mode=TreeDictIteratorMode.KEY_AND_VALUE)

    def value_at_index(self, idx: int) -> Any:
        """Return the value at the index.

        Args:
            idx: The index.

        Returns:
            The value of the node at the index.
        """
        return self._node_at_index(idx).value

    def key_and_value_at_index(self, idx: int) -> tuple[Any, Any]:
        """Return the key-value pair at the index.

        Args:
            idx: The index.

        Returns:
            The key-value pair of the node at the index.
        """
        node = self._node_at_index(idx)

        return (node.key, node.value)

    def add(self, key: Any, value: Any) -> None:
        """Insert a key and value.

        Args:
            key: The key to be added.
            value: The associated value.
        """
        self.insert(key, value)

    def insert(self, key: Any, value: Any) -> None:
        """Insert a key and value.

        Args:
            key: The key to be added.
            value: The associated value.

        Raises:
            KeyError: If the key already exists.
        """
        self._temp_attributes = (key, value)
        self._insert_key(key)
        self._temp_attributes = None

__getitem__(key: Any) -> Any

Return the value associated with the key.

Overrides the method of the base class where the key at a specified index is returned.

Parameters:

Name	Type	Description	Default
key	Any	The key.	required

Returns:

Type	Description
Any	The value associated with the key.

Raises:

Type	Description
KeyError	If the key does not exist.

Source code in tralda/datastructures/bst/avl.py

def __getitem__(self, key: Any) -> Any:
    """Return the value associated with the key.

    Overrides the method of the base class where the key at a specified index is returned.

    Args:
        key: The key.

    Returns:
        The value associated with the key.

    Raises:
        KeyError: If the key does not exist.
    """
    node = self._find(key)

    if not node:
        raise KeyError(f"key {key} is not in the TreeDict")

    return node.value

get(key: Any, default: Any = None) -> Any

Return the value associated with the key or a default value if the key does not exist.

Parameters:

Name	Type	Description	Default
key	Any	The key for which to get the associated value.	required
default	Any	The default to be returned if the key does not exist.	None

Returns:

Type	Description
Any	The value associated with the key or the default value.

Source code in tralda/datastructures/bst/avl.py

def get(self, key: Any, default: Any = None) -> Any:
    """Return the value associated with the key or a default value if the key does not exist.

    Args:
        key: The key for which to get the associated value.
        default: The default to be returned if the key does not exist.

    Returns:
        The value associated with the key or the default value.
    """
    node = self._find(key)

    if node:
        return node.value
    else:
        return default

keys() -> Iterator[Any]

Iterator for the keys.

Returns:

Type	Description
Iterator[Any]	An iterator for the keys.

Source code in tralda/datastructures/bst/avl.py

def keys(self) -> Iterator[Any]:
    """Iterator for the keys.

    Returns:
        An iterator for the keys.
    """
    return self.iterator_class(self, mode=TreeDictIteratorMode.KEY)

values() -> Iterator[Any]

Iterator for the values.

Returns:

Type	Description
Iterator[Any]	An iterator for the values.

Source code in tralda/datastructures/bst/avl.py

def values(self) -> Iterator[Any]:
    """Iterator for the values.

    Returns:
        An iterator for the values.
    """
    return self.iterator_class(self, mode=TreeDictIteratorMode.VALUE)

items() -> Iterator[tuple[Any, Any]]

Iterator for the key-value pairs.

Returns:

Type	Description
Iterator[tuple[Any, Any]]	An iterator for the key-value pairs.

Source code in tralda/datastructures/bst/avl.py

def items(self) -> Iterator[tuple[Any, Any]]:
    """Iterator for the key-value pairs.

    Returns:
        An iterator for the key-value pairs.
    """
    return self.iterator_class(self, mode=TreeDictIteratorMode.KEY_AND_VALUE)

value_at_index(idx: int) -> Any

Return the value at the index.

Parameters:

Name	Type	Description	Default
idx	int	The index.	required

Returns:

Type	Description
Any	The value of the node at the index.

Source code in tralda/datastructures/bst/avl.py

def value_at_index(self, idx: int) -> Any:
    """Return the value at the index.

    Args:
        idx: The index.

    Returns:
        The value of the node at the index.
    """
    return self._node_at_index(idx).value

key_and_value_at_index(idx: int) -> tuple[Any, Any]

Return the key-value pair at the index.

Parameters:

Name	Type	Description	Default
idx	int	The index.	required

Returns:

Type	Description
tuple[Any, Any]	The key-value pair of the node at the index.

Source code in tralda/datastructures/bst/avl.py

def key_and_value_at_index(self, idx: int) -> tuple[Any, Any]:
    """Return the key-value pair at the index.

    Args:
        idx: The index.

    Returns:
        The key-value pair of the node at the index.
    """
    node = self._node_at_index(idx)

    return (node.key, node.value)

add(key: Any, value: Any) -> None

Insert a key and value.

Parameters:

Name	Type	Description	Default
key	Any	The key to be added.	required
value	Any	The associated value.	required

Source code in tralda/datastructures/bst/avl.py

def add(self, key: Any, value: Any) -> None:
    """Insert a key and value.

    Args:
        key: The key to be added.
        value: The associated value.
    """
    self.insert(key, value)

insert(key: Any, value: Any) -> None

Insert a key and value.

Parameters:

Name	Type	Description	Default
key	Any	The key to be added.	required
value	Any	The associated value.	required

Raises:

Type	Description
KeyError	If the key already exists.

Source code in tralda/datastructures/bst/avl.py

def insert(self, key: Any, value: Any) -> None:
    """Insert a key and value.

    Args:
        key: The key to be added.
        value: The associated value.

    Raises:
        KeyError: If the key already exists.
    """
    self._temp_attributes = (key, value)
    self._insert_key(key)
    self._temp_attributes = None

---

Red-Black-tree-based data structures

TreeSet

Bases: BaseBinarySearchTree

Red-black tree.

Source code in tralda/datastructures/bst/red_black.py

class TreeSet(BaseBinarySearchTree):
    """Red-black tree."""

    __slots__ = ()

    node_class = RedBlackTreeNode
    iterator_class: Iterator[Any] = BinaryTreeIterator

    def __init__(self) -> None:
        """Constructor of the TreeSet class (red-black tree implementation)."""
        super().__init__()

    @classmethod
    def join(cls, tree_left: TreeSet, tree_right: TreeSet, key: Any | None = None) -> TreeSet:
        """Join function for two red-black trees.

        Join two red-black trees T_L and T_R with x < key < y for all x in T_L and y in T_R. The
        resulting tree will contain all keys in T_L and T_R and additionally the new provided key.
        If no key is provided, a dummy key will be added during the join operation, which will be
        removed in the end, keeping the total O(log(n)) time complexity.

        NOTE: The input red-black tree instances should no longer be used afterwards.

        Args:
            tree_left: The tree T_L with the smaller elements.
            tree_right: The tree T_R with the larger elements.
            key: The new key to be added. If None, a dummy node will be temporarily inserted and
                removed in the end.

        Raises:
            TypeError: If one of the two trees is not an instance of the correct red-black tree
                type.

        Returns:
            The joined tree.

        References:
            1. https://en.wikipedia.org/wiki/Red-black_tree#Set_operations_and_bulk_operations
        """
        if not isinstance(tree_left, cls):
            raise TypeError(f"tree_left must be of type {cls}, but is {type(tree_left)}")
        if not isinstance(tree_right, cls):
            raise TypeError(f"tree_right must be of type {cls}, but is {type(tree_right)}")

        new_node = cls.node_class(0) if key is None else cls.node_class(key)

        tree = cls._join(tree_left, tree_right, new_node)

        # if no key was provided, remove the dummy node
        if key is None:
            tree._delete_node(new_node)

        return tree

    def split_at_node(
        self,
        node: RedBlackTreeNode,
        keep_node_left: bool = False,
        keep_node_right: bool = False,
    ) -> tuple[TreeSet, TreeSet]:
        """Split the tree at a given node.

        Considering the key of the provided node, the items x in the first set (left tree) will
        satisfy x <= key, and the items x in the second set (right tree) will satisfy x >= key.
        The key itself is only optionally retained in one of the two sets.

        NOTE: The red-black tree instance should no longer be used afterwards.

        Args:
            node: The node at which to split the tree.
            keep_node_left: Whether to keep the provided node in the set with the smaller keys
                (left tree).
            keep_node_right: Whether to keep the provided node in the set with the larger keys
                (right tree).

        Raises:
            ValueError: If both keep_node_left and keep_node_right are set to True.

        Returns:
            A tuple containing the left tree and the right tree (i.e., the sets after splitting).

        References:
            1. https://en.wikipedia.org/wiki/Red-black_tree#Set_operations_and_bulk_operations
        """
        if keep_node_left and keep_node_right:
            raise ValueError("node cannot be kept in both left are right tree")

        return self._split_at_node(node, keep_node_left, keep_node_right)

    def check_integrity(self, verbose: bool = False) -> bool:
        """Integrity check of the tree.

        Checks whether the size and height is correct in all subtrees. Moreover, the red-black
        properties are checked. Intended for debugging and testing purpose.

        Args:
            verbose: If True print where the integrity check has failed.

        Returns:
            Whether all integrity checks have been passed.
        """
        # check that all parents are set correctly
        for i, node in enumerate(self._postorder_traversal()):
            if node.left and node.left.parent is not node:
                if verbose:
                    print(f"node {node.left} should have parent {node} but has {node.left.parent}")
                return False

            if node.right and node.right.parent is not node:
                if verbose:
                    print(
                        f"node {node.right} should have parent {node} but has {node.right.parent}"
                    )
                return False

        # check size and height
        if not super().check_integrity(verbose=verbose):
            return False

        # check the red-black properties
        for i, node in enumerate(self._postorder_traversal()):
            # (1) all patch from the root to a leaf (counting NIL nodes as black leaves) have the
            # same number of black nodes
            black_height_left = node.left.black_height if node.left else 1
            black_height_right = node.right.black_height if node.right else 1

            if black_height_left != black_height_right:
                if verbose:
                    print(
                        f"children of node {node} have different black height "
                        f"({black_height_left} vs. {black_height_right}, traversal step: {i})"
                    )
                return False

            expected_black_height = black_height_left + int(node.is_black)

            if expected_black_height != node.black_height:
                if verbose:
                    print(
                        f"node {node} should black height {expected_black_height} but has "
                        f"{node.black_height} (traversal step: {i})"
                    )
                return False

            # (2) a red node does not have a red child
            if node.is_red and node.parent and node.parent.is_red:
                if verbose:
                    print(f"node {node.parent} is red and has a red child {node}")
                return False

        return True

    def _copy_subtree(self, node: RedBlackTreeNode) -> RedBlackTreeNode:
        """Recursive auxiliary function to copy a subtree under the provided node.

        Args:
            node: Node whose subtree to copy.

        Returns:
            The root node of the copied subtree.
        """
        node_copy = node.copy()
        if node.left:
            node_copy.left = self._copy_subtree(node.left)
            node_copy.left.parent = node_copy
        if node.right:
            node_copy.right = self._copy_subtree(node.right)
            node_copy.right.parent = node_copy

        return node_copy

    def _traverse_to_root_and_update(self, node: RedBlackTreeNode) -> None:
        """Update all nodes on the path to the root.

        Args:
            node: Node for starting the traversal to the root.
        """
        while node:
            node.update()
            node = node.parent

    def _replace_child(
        self,
        parent: RedBlackTreeNode | None,
        old_child: RedBlackTreeNode,
        new_child: RedBlackTreeNode | None,
    ) -> None:
        """Replace the child of a parent node.

        Args:
            parent: The parent node.
            old_child: The current child node to be replaced.
            new_child: The replacement node.
        """
        if not parent:
            self.root = new_child
        elif parent.right is old_child:
            parent.right = new_child
        elif parent.left is old_child:
            parent.left = new_child
        else:
            raise RuntimeError("node is not a child of the provided parent")

        if new_child:
            new_child.parent = parent

    def _rotate_subtree(self, node: RedBlackTreeNode, direction: int) -> RedBlackTreeNode:
        """Rotate the node and and one of its children.

        Args:
            node: The node that will become a child of one of its children.
            direction: The direction of the rotation (0 = left rotation, 1 = right rotation).

        Returns:
            The child of node that is now its parent, i.e., the new root of the subtree.
        """
        parent = node.parent
        new_root = node.child(1 - direction)
        new_child = new_root.child(direction)

        node.set_child(1 - direction, new_child)

        if new_child:
            new_child.parent = node

        new_root.set_child(direction, node)
        node.parent = new_root
        self._replace_child(parent, node, new_root)

        # node may no longer be on the path from the inserted / moved up node to the root and,
        # therefore, it must be updated here already
        node.update()

        return new_root

    def _insert_key(self, key: Any) -> None:
        """Insert a key into the tree if not already present.

        Args:
            key: The key to be inserted.

        Raises:
            KeyError: If the key already exists.
        """
        node = self.root
        parent = None
        direction = 0

        # search for the place to insert the new key
        while node:
            parent = node
            if key < node.key:
                node = node.left
                direction = 0
            elif key > node.key:
                node = node.right
                direction = 1
            else:
                raise KeyError(f"key {key} already exists")

        # create and insert the new node, rebalance, update nodes on the way to the root
        new_node = self.node_class(key)
        self._insert_node(new_node, parent, direction)

    def _insert_node(
        self,
        node: RedBlackTreeNode,
        parent: RedBlackTreeNode | None,
        direction: int,
    ):
        """Insert a node at the specified location, rebalance, and update the nodes.

        Args:
            node: The new node to insert.
            parent: The parent below which to insert the new node.
            direction: The direction where to insert the new node below parent (0 = left, 1 =
                right).
        """
        self._insert_node_and_rebalance(node, parent, direction)
        self._traverse_to_root_and_update(node)

    def _insert_node_and_rebalance(
        self,
        node: RedBlackTreeNode,
        parent: RedBlackTreeNode | None,
        direction: int,
    ) -> None:
        """Insert a node at the specified location.

        Args:
            node: The new node to insert.
            parent: The parent below which to insert the new node.
            direction: The direction where to insert the new node below parent (0 = left, 1 =
                right).

        References:
            1. https://en.wikipedia.org/wiki/Red-black_tree#Insertion
        """
        # cases follow the implementation in https://en.wikipedia.org/wiki/Red-black_tree#Insertion
        # (accessed Nov 29, 2025)

        node.turn_red()
        node.parent = parent

        # tree was empty before --> new node becomes the root
        if not parent:
            self.root = node
            return

        parent.set_child(direction, node)

        # rebalance the tree
        while parent:
            # case 1
            if parent.is_black:
                return

            grandparent = parent.parent

            if not grandparent:
                # case 4
                parent.turn_black()
                return

            direction = parent.direction
            uncle = grandparent.child(1 - direction)

            if not uncle or uncle.is_black:
                if node is parent.child(1 - direction):
                    # case 5
                    self._rotate_subtree(parent, direction)
                    node = parent
                    parent = grandparent.child(direction)

                # case 6
                self._rotate_subtree(grandparent, 1 - direction)
                parent.turn_black()
                grandparent.turn_red()
                return

            # case 2
            parent.turn_black()
            uncle.turn_black()
            grandparent.turn_red()
            node = grandparent

            parent = node.parent

    def _delete_key(self, key: Any) -> None:
        """Delete a key.

        Args:
            key: The key to be deleted.
        """
        node = self._find(key)

        if not node:
            raise KeyError(f"key {key} not found")

        self._delete_node(node)

    def _pop_at_index(self, idx: int) -> None:
        """Remove item at the index and return it.

        Args:
            idx: The index of the element to be removed.
        """
        node = self._node_at_index(idx)
        self._temp_attributes = node.get_attributes()
        self._delete_node(node)

    def _replace_node_to_remove(self, node: RedBlackTreeNode) -> RedBlackTreeNode:
        """Find the inorder successor of a node to be removed instead.

        The inorder successor is guaranteed to have at most one child (the right child). The
        attributes of this node will be copied into the original node to remove.

        Args:
            node: The node whose content shall be removed.

        Returns:
            The instance of the inorder successor node which can now be removed instead.
        """
        # find smallest node of right subtree ("inorder successor" of current node)
        inorder_successor = node.right.smallest_in_subtree()

        # copy inorder successor's data to current node (but keep its color)
        inorder_successor.copy_attributes_to_node(node)

        return inorder_successor

    def _delete_node(self, node: RedBlackTreeNode) -> None:
        """Delete a node from the tree.

        Args:
            node: The node to delete from the tree.

        References:
            1. https://en.wikipedia.org/wiki/Red-black_tree#Removal
        """
        # node to delete has 2 children
        if node.left and node.right:
            node = self._replace_node_to_remove(node)

        parent = node.parent

        # node to delete has 1 child (single child must be red, deleted node must be black)
        if node.left or node.right:
            self._delete_node_with_one_child(node)
        # node to delete has no children and is the root
        elif not parent:
            self.root = None  # the tree is now empty
            return
        # node to delete has no children and is red --> simply delete the leaf
        elif node.is_red:
            self._replace_child(parent, node, None)
        # node to delete has no children and black --> deleting creates imbalance and requires
        # rebalancing
        else:
            self._delete_black_leaf(node)

        self._traverse_to_root_and_update(parent)

    def _delete_node_with_one_child(self, node: RedBlackTreeNode) -> None:
        """Delete with exactly one child from the tree.

        Args:
            node: The node to delete from the tree.
        """
        if node.left:
            self._replace_child(node.parent, node, node.left)
            node.left.turn_black()
        else:
            self._replace_child(node.parent, node, node.right)
            node.right.turn_black()

    def _delete_black_leaf(self, node: RedBlackTreeNode) -> None:
        """Remove a black leaf from the tree.

        Args:
            node: The black leaf node to be removed.
        """
        # cases follow the implementation in https://en.wikipedia.org/wiki/Red–black_tree#Removal
        # (accessed Nov 29, 2025)

        parent = node.parent
        sibling = None
        close_nephew = None
        distant_nephew = None
        direction = node.direction
        skip_case_5 = False

        parent.set_child(direction, None)

        while True:
            sibling = parent.child(1 - direction)
            distant_nephew = sibling.child(1 - direction)
            close_nephew = sibling.child(direction)

            if sibling.is_red:
                # case 3
                self._rotate_subtree(parent, direction)
                parent.turn_red()
                sibling.turn_black()
                sibling = close_nephew

                distant_nephew = sibling.child(1 - direction)
                if distant_nephew and distant_nephew.is_red:
                    skip_case_5 = True
                    break  # go to case 6
                close_nephew = sibling.child(direction)

                if close_nephew and close_nephew.is_red:
                    break  # go to case 5

                # case 4
                sibling.turn_red()
                parent.turn_black()
                return

            if distant_nephew and distant_nephew.is_red:
                skip_case_5 = True
                break  # go to case 6

            if close_nephew and close_nephew.is_red:
                break  # go to case 5

            if not parent:
                # case 1
                return

            if parent.is_red:
                # case 4
                sibling.turn_red()
                parent.turn_black()
                return

            # case 2
            sibling.turn_red()
            node = parent

            parent = node.parent
            if not parent:
                return

            direction = node.direction

        if not skip_case_5:
            # case 5
            self._rotate_subtree(sibling, 1 - direction)
            sibling.turn_red()
            close_nephew.turn_black()
            distant_nephew = sibling
            sibling = close_nephew

        # case 6
        self._rotate_subtree(parent, direction)
        if parent.is_red:
            sibling.turn_red()
        else:
            sibling.turn_black()
        parent.turn_black()
        distant_nephew.turn_black()

    @classmethod
    def _join(cls, tree_left: TreeSet, tree_right: TreeSet, new_node: RedBlackTreeNode) -> TreeSet:
        """Join function for two red-black trees.

        Join two red-black trees T_L and T_R using a new node. The resulting tree will contain all
        keys in T_L, the new key, and all keys in T_R (order w.r.t. to an in-order traversal).

        Args:
            tree_left: The tree T_L.
            tree_right: The tree T_R.
            key: The new node the will be between the nodes of T_L and T_R in an in-order traversal.

        Returns:
            The joined tree.
        """
        # both trees are empty --> reuse instance tree_left
        if not tree_left and not tree_right:
            tree_left.root = new_node
            return tree_left
        # only tree_right is non-empty --> insert new_node as new smallest element
        elif not tree_left:
            parent = tree_right.root.smallest_in_subtree()
            tree_right._insert_node(new_node, parent, 0)
            return tree_right
        # only tree_left is non-empty --> insert new_node as new largest element
        elif not tree_right:
            parent = tree_left.root.largest_in_subtree()
            tree_left._insert_node(new_node, parent, 1)
            return tree_left

        # create a new tree if both trees are non-empty
        tree = cls()
        tree.root = tree._join_nonempty_trees(tree_left.root, tree_right.root, new_node)

        # update all nodes on the path from the new node to the root
        tree._traverse_to_root_and_update(new_node)

        return tree

    def _join_subtrees(
        self,
        parent: RedBlackTreeNode,
        left_child: RedBlackTreeNode | None,
        right_child: RedBlackTreeNode | None,
    ) -> None:
        """Attach two nodes to a common parent.

        Args:
            parent: The parent node.
            left_child: The new left child.
            right_child: The new right child.
        """
        parent.left = left_child
        if left_child:
            left_child.parent = parent

        parent.right = right_child
        if right_child:
            right_child.parent = parent

    def _join_nonempty_trees(
        self,
        root_left: RedBlackTreeNode,
        root_right: RedBlackTreeNode,
        new_node: RedBlackTreeNode,
    ) -> RedBlackTreeNode:
        """Join two non-empty red-black trees.

        Args:
            root_left: The root of the left tree (containing the smaller elements).
            root_right: The root of the right tree (containing the larger elements).
            new_node: The new node used to join the trees.

        Returns:
            The root of the joined tree.
        """
        if root_left.black_height > root_right.black_height:
            new_root = self._join_right(root_left, root_right, new_node)
            new_root.turn_black()
        elif root_left.black_height < root_right.black_height:
            new_root = self._join_left(root_left, root_right, new_node)
            new_root.turn_black()
        else:
            new_root = new_node
            self._join_subtrees(new_root, root_left, root_right)

        return new_root

    def _join_right(
        self,
        root_left: RedBlackTreeNode,
        root_right: RedBlackTreeNode,
        new_node: RedBlackTreeNode,
    ) -> RedBlackTreeNode:
        """Recursively move done the right spine of the left tree to join the trees.

        Args:
            root_left: The root of the left tree (containing the smaller elements).
            root_right: The root of the right tree (containing the larger elements).
            new_node: The new node used to join the trees.

        Returns:
            The root of the joined tree.
        """
        if not root_left or (
            root_left.is_black and root_left.black_height == root_right.black_height
        ):
            new_node.turn_red()
            self._join_subtrees(new_node, root_left, root_right)
            return new_node

        root_left.right = self._join_right(root_left.right, root_right, new_node)
        root_left.right.parent = root_left

        if root_left.is_black and root_left.right.is_red and root_left.right.right.is_red:
            root_left.right.right.turn_black()
            return self._rotate_subtree(root_left, 0)  # left rotation

        return root_left

    def _join_left(
        self,
        root_left: RedBlackTreeNode,
        root_right: RedBlackTreeNode,
        new_node: RedBlackTreeNode,
    ) -> RedBlackTreeNode:
        """Recursively move done the left spine of the right tree to join the trees.

        Args:
            root_left: The root of the left tree (containing the smaller elements).
            root_right: The root of the right tree (containing the larger elements).
            new_node: The new node used to join the trees.

        Returns:
            The root of the joined tree.
        """
        if not root_right or (
            root_right.is_black and root_left.black_height == root_right.black_height
        ):
            new_node.turn_red()
            self._join_subtrees(new_node, root_left, root_right)
            return new_node

        root_right.left = self._join_left(root_left, root_right.left, new_node)
        root_right.left.parent = root_right

        if root_right.is_black and root_right.left.is_red and root_right.left.left.is_red:
            root_right.left.left.turn_black()
            return self._rotate_subtree(root_right, 1)  # right rotation

        return root_right

    def _split_at_node(
        self,
        node: RedBlackTreeNode,
        keep_node_left: bool,
        keep_node_right: bool,
    ) -> tuple[TreeSet, TreeSet]:
        """Split the tree at a given node.

        Args:
            node: The node at which to split the tree.
            keep_node_left: Whether to keep the provided node in the left tree.
            keep_node_right: Whether to keep the provided node in the right tree.

        Returns:
            A tuple containing the left tree and the right tree (i.e., the sets after splitting).
        """
        cls = self.__class__
        tree_left = cls()
        tree_right = cls()

        # initialize the two trees with the subtrees of the node's children
        tree_left.root = node.left
        tree_right.root = node.right
        node.detach_children()
        node.turn_black()

        parent = node.parent
        if parent:
            direction_parent = node.direction
            parent.set_child(direction_parent, None)
            node.parent = None

        # keep the node in the left tree
        if keep_node_left and not tree_left:
            tree_left.root = node
        elif keep_node_left:
            rightmost_node = tree_left.root.largest_in_subtree()
            tree_left._insert_node(node, rightmost_node, 1)

        # keep the node in the right tree
        if keep_node_right and not tree_right:
            tree_right.root = node
        elif keep_node_right:
            leftmost_node = tree_right.root.smallest_in_subtree()
            tree_right._insert_node(node, leftmost_node, 0)

        # now traverse to the root
        while parent:
            node = parent
            direction = direction_parent
            parent = node.parent
            if parent:
                direction_parent = node.direction
                parent.set_child(direction_parent, None)
                node.parent = None

            other_child = node.child(1 - direction)
            node.detach_children()
            node.turn_black()

            # create a subtree rooted at other_child that can be joined with tree_left / tree_right
            subtree = cls()
            subtree.root = other_child

            # coming from left --> node and subtree have to go into the right subtree
            if direction == 0:
                tree_right = cls._join(tree_right, subtree, node)
            # coming from right --> node and subtree have to go into the left subtree
            else:
                tree_left = cls._join(subtree, tree_left, node)

        return tree_left, tree_right

__init__() -> None

Constructor of the TreeSet class (red-black tree implementation).

Source code in tralda/datastructures/bst/red_black.py

def __init__(self) -> None:
    """Constructor of the TreeSet class (red-black tree implementation)."""
    super().__init__()

join(tree_left: TreeSet, tree_right: TreeSet, key: Any | None = None) -> TreeSet classmethod

Join function for two red-black trees.

Join two red-black trees T_L and T_R with x < key < y for all x in T_L and y in T_R. The resulting tree will contain all keys in T_L and T_R and additionally the new provided key. If no key is provided, a dummy key will be added during the join operation, which will be removed in the end, keeping the total O(log(n)) time complexity.

NOTE: The input red-black tree instances should no longer be used afterwards.

Parameters:

Name	Type	Description	Default
tree_left	TreeSet	The tree T_L with the smaller elements.	required
tree_right	TreeSet	The tree T_R with the larger elements.	required
key	Any | None	The new key to be added. If None, a dummy node will be temporarily inserted and removed in the end.	None

Raises:

Type	Description
TypeError	If one of the two trees is not an instance of the correct red-black tree type.

Returns:

Type	Description
TreeSet	The joined tree.

References

1. https://en.wikipedia.org/wiki/Red-black_tree#Set_operations_and_bulk_operations

Source code in tralda/datastructures/bst/red_black.py

@classmethod
def join(cls, tree_left: TreeSet, tree_right: TreeSet, key: Any | None = None) -> TreeSet:
    """Join function for two red-black trees.

    Join two red-black trees T_L and T_R with x < key < y for all x in T_L and y in T_R. The
    resulting tree will contain all keys in T_L and T_R and additionally the new provided key.
    If no key is provided, a dummy key will be added during the join operation, which will be
    removed in the end, keeping the total O(log(n)) time complexity.

    NOTE: The input red-black tree instances should no longer be used afterwards.

    Args:
        tree_left: The tree T_L with the smaller elements.
        tree_right: The tree T_R with the larger elements.
        key: The new key to be added. If None, a dummy node will be temporarily inserted and
            removed in the end.

    Raises:
        TypeError: If one of the two trees is not an instance of the correct red-black tree
            type.

    Returns:
        The joined tree.

    References:
        1. https://en.wikipedia.org/wiki/Red-black_tree#Set_operations_and_bulk_operations
    """
    if not isinstance(tree_left, cls):
        raise TypeError(f"tree_left must be of type {cls}, but is {type(tree_left)}")
    if not isinstance(tree_right, cls):
        raise TypeError(f"tree_right must be of type {cls}, but is {type(tree_right)}")

    new_node = cls.node_class(0) if key is None else cls.node_class(key)

    tree = cls._join(tree_left, tree_right, new_node)

    # if no key was provided, remove the dummy node
    if key is None:
        tree._delete_node(new_node)

    return tree

split_at_node(node: RedBlackTreeNode, keep_node_left: bool = False, keep_node_right: bool = False) -> tuple[TreeSet, TreeSet]

Split the tree at a given node.

Considering the key of the provided node, the items x in the first set (left tree) will satisfy x <= key, and the items x in the second set (right tree) will satisfy x >= key. The key itself is only optionally retained in one of the two sets.

NOTE: The red-black tree instance should no longer be used afterwards.

Parameters:

Name	Type	Description	Default
node	RedBlackTreeNode	The node at which to split the tree.	required
keep_node_left	bool	Whether to keep the provided node in the set with the smaller keys (left tree).	False
keep_node_right	bool	Whether to keep the provided node in the set with the larger keys (right tree).	False

Raises:

Type	Description
ValueError	If both keep_node_left and keep_node_right are set to True.

Returns:

Type	Description
tuple[TreeSet, TreeSet]	A tuple containing the left tree and the right tree (i.e., the sets after splitting).

References

1. https://en.wikipedia.org/wiki/Red-black_tree#Set_operations_and_bulk_operations

Source code in tralda/datastructures/bst/red_black.py

def split_at_node(
    self,
    node: RedBlackTreeNode,
    keep_node_left: bool = False,
    keep_node_right: bool = False,
) -> tuple[TreeSet, TreeSet]:
    """Split the tree at a given node.

    Considering the key of the provided node, the items x in the first set (left tree) will
    satisfy x <= key, and the items x in the second set (right tree) will satisfy x >= key.
    The key itself is only optionally retained in one of the two sets.

    NOTE: The red-black tree instance should no longer be used afterwards.

    Args:
        node: The node at which to split the tree.
        keep_node_left: Whether to keep the provided node in the set with the smaller keys
            (left tree).
        keep_node_right: Whether to keep the provided node in the set with the larger keys
            (right tree).

    Raises:
        ValueError: If both keep_node_left and keep_node_right are set to True.

    Returns:
        A tuple containing the left tree and the right tree (i.e., the sets after splitting).

    References:
        1. https://en.wikipedia.org/wiki/Red-black_tree#Set_operations_and_bulk_operations
    """
    if keep_node_left and keep_node_right:
        raise ValueError("node cannot be kept in both left are right tree")

    return self._split_at_node(node, keep_node_left, keep_node_right)

check_integrity(verbose: bool = False) -> bool

Integrity check of the tree.

Checks whether the size and height is correct in all subtrees. Moreover, the red-black properties are checked. Intended for debugging and testing purpose.

Parameters:

Name	Type	Description	Default
verbose	bool	If True print where the integrity check has failed.	False

Returns:

Type	Description
bool	Whether all integrity checks have been passed.

Source code in tralda/datastructures/bst/red_black.py

def check_integrity(self, verbose: bool = False) -> bool:
    """Integrity check of the tree.

    Checks whether the size and height is correct in all subtrees. Moreover, the red-black
    properties are checked. Intended for debugging and testing purpose.

    Args:
        verbose: If True print where the integrity check has failed.

    Returns:
        Whether all integrity checks have been passed.
    """
    # check that all parents are set correctly
    for i, node in enumerate(self._postorder_traversal()):
        if node.left and node.left.parent is not node:
            if verbose:
                print(f"node {node.left} should have parent {node} but has {node.left.parent}")
            return False

        if node.right and node.right.parent is not node:
            if verbose:
                print(
                    f"node {node.right} should have parent {node} but has {node.right.parent}"
                )
            return False

    # check size and height
    if not super().check_integrity(verbose=verbose):
        return False

    # check the red-black properties
    for i, node in enumerate(self._postorder_traversal()):
        # (1) all patch from the root to a leaf (counting NIL nodes as black leaves) have the
        # same number of black nodes
        black_height_left = node.left.black_height if node.left else 1
        black_height_right = node.right.black_height if node.right else 1

        if black_height_left != black_height_right:
            if verbose:
                print(
                    f"children of node {node} have different black height "
                    f"({black_height_left} vs. {black_height_right}, traversal step: {i})"
                )
            return False

        expected_black_height = black_height_left + int(node.is_black)

        if expected_black_height != node.black_height:
            if verbose:
                print(
                    f"node {node} should black height {expected_black_height} but has "
                    f"{node.black_height} (traversal step: {i})"
                )
            return False

        # (2) a red node does not have a red child
        if node.is_red and node.parent and node.parent.is_red:
            if verbose:
                print(f"node {node.parent} is red and has a red child {node}")
            return False

    return True