pkstruct 0.1.0__py3-none-any.whl

pkstruct/trees/balancing.py

@@ -0,0 +1,253 @@
+"""
+pkstruct.trees.balancing
+========================
+Shared balancing utilities used by :class:`~pkstruct.trees.avl.AVLTree`
+and :class:`~pkstruct.trees.red_black.RedBlackTree`.
+
+Design goals
+------------
+* **Single** ``rotate`` function accepts a ``direction`` parameter instead
+  of separate ``rotate_left`` / ``rotate_right`` methods.
+* ``update_metadata`` is dispatched on node type so AVL and RB nodes each
+  maintain their own invariants without branching inside tree classes.
+* All functions operate on nodes directly and return the new sub-tree root
+  when the root may change; callers are responsible for re-attaching.
+"""
+
+from __future__ import annotations
+
+from .node import AVLNode, RBNode, TreeNode
+
+__all__ = [
+    "rotate",
+    "rebalance",
+    "update_metadata",
+    "get_balance_factor",
+    "validate_balance",
+]
+
+_AnyNode = AVLNode | RBNode | TreeNode
+
+
+# ---------------------------------------------------------------------------
+# Height helpers (AVL)
+# ---------------------------------------------------------------------------
+
+
+def _height(node: AVLNode | None) -> int:
+    """Return the height of *node*, or 0 if *None*."""
+    return node.height if node is not None else 0
+
+
+def _update_height(node: AVLNode) -> None:
+    """Recalculate ``node.height`` from its children."""
+    node.height = 1 + max(_height(node.left), _height(node.right))  # type: ignore[arg-type]
+
+
+# ---------------------------------------------------------------------------
+# Core rotation
+# ---------------------------------------------------------------------------
+
+
+def rotate(
+    node: _AnyNode,
+    direction: str,
+    root: _AnyNode | None = None,
+) -> _AnyNode:
+    """
+    Perform a single rotation on *node* and return the new sub-tree root.
+
+    Parameters
+    ----------
+    node:
+        The node to rotate around (the "pivot's parent").
+    direction:
+        ``"left"`` or ``"right"``.
+    root:
+        The overall tree root.  When *node* **is** the tree root the
+        caller must replace ``tree.root`` with the returned value.
+
+    Returns
+    -------
+    _AnyNode
+        The new sub-tree root after rotation.
+
+    Raises
+    ------
+    ValueError
+        If *direction* is not ``"left"`` or ``"right"``.
+    ValueError
+        If the required child is absent (cannot rotate).
+    """
+    if direction not in ("left", "right"):
+        raise ValueError(f"direction must be 'left' or 'right', got {direction!r}.")
+
+    if direction == "left":
+        pivot = node.right
+        if pivot is None:
+            raise ValueError("Cannot rotate left: node has no right child.")
+        node.right = pivot.left
+        if pivot.left is not None and pivot.left.key is not None:
+            pivot.left.parent = node
+        pivot.parent = node.parent
+        if node.parent is None:
+            pass  # caller updates root
+        elif node.parent.left is node:
+            node.parent.left = pivot
+        else:
+            node.parent.right = pivot
+        pivot.left = node
+        node.parent = pivot
+    else:  # right
+        pivot = node.left
+        if pivot is None:
+            raise ValueError("Cannot rotate right: node has no left child.")
+        node.left = pivot.right
+        if pivot.right is not None and pivot.right.key is not None:
+            pivot.right.parent = node
+        pivot.parent = node.parent
+        if node.parent is None:
+            pass
+        elif node.parent.left is node:
+            node.parent.left = pivot
+        else:
+            node.parent.right = pivot
+        pivot.right = node
+        node.parent = pivot
+
+    # Update heights for AVL nodes
+    if isinstance(node, AVLNode):
+        _update_height(node)
+        _update_height(pivot)  # type: ignore[arg-type]
+
+    return pivot  # type: ignore[return-value]
+
+
+# ---------------------------------------------------------------------------
+# Metadata update dispatch
+# ---------------------------------------------------------------------------
+
+
+def update_metadata(node: _AnyNode) -> None:
+    """
+    Recompute cached metadata stored on *node*.
+
+    For :class:`~pkstruct.trees.node.AVLNode`: updates ``height``.
+    For other node types: no-op (override as needed).
+
+    Parameters
+    ----------
+    node:
+        Node whose metadata should be refreshed.
+    """
+    if isinstance(node, AVLNode):
+        _update_height(node)
+
+
+# ---------------------------------------------------------------------------
+# Balance factor
+# ---------------------------------------------------------------------------
+
+
+def get_balance_factor(node: AVLNode | None) -> int:  # type: ignore[return]
+    """
+    Return the balance factor of *node* (left height − right height).
+
+    Parameters
+    ----------
+    node:
+        An :class:`~pkstruct.trees.node.AVLNode`, or *None*.
+
+    Returns
+    -------
+    int
+        Balance factor; **0** when *node* is *None*.
+    """
+    if node is None:
+        return 0
+    return _height(node.left) - _height(node.right)  # type: ignore[arg-type]
+
+
+# ---------------------------------------------------------------------------
+# AVL rebalance
+# ---------------------------------------------------------------------------
+
+
+def rebalance(node: AVLNode, strategy: str = "avl") -> AVLNode:
+    """
+    Restore balance at *node* using the requested *strategy*.
+
+    Currently the only supported strategy is ``"avl"``.
+
+    Parameters
+    ----------
+    node:
+        Node to rebalance.
+    strategy:
+        Balancing strategy identifier (currently only ``"avl"``).
+
+    Returns
+    -------
+    AVLNode
+        The new sub-tree root after rebalancing (may differ from *node*).
+
+    Raises
+    ------
+    ValueError
+        If *strategy* is not recognised.
+    """
+    if strategy != "avl":
+        raise ValueError(f"Unknown rebalancing strategy: {strategy!r}.")
+
+    _update_height(node)
+    balance = get_balance_factor(node)
+
+    # Left heavy
+    if balance > 1:
+        left_child: AVLNode = node.left  # type: ignore[assignment]
+        if get_balance_factor(left_child) < 0:
+            # Left-Right case
+            node.left = rotate(left_child, "left")  # type: ignore[assignment]
+            node.left.parent = node
+        return rotate(node, "right")  # type: ignore[return-value]
+
+    # Right heavy
+    if balance < -1:
+        right_child: AVLNode = node.right  # type: ignore[assignment]
+        if get_balance_factor(right_child) > 0:
+            # Right-Left case
+            node.right = rotate(right_child, "right")  # type: ignore[assignment]
+            node.right.parent = node
+        return rotate(node, "left")  # type: ignore[return-value]
+
+    return node
+
+
+# ---------------------------------------------------------------------------
+# Validation
+# ---------------------------------------------------------------------------
+
+
+def validate_balance(root: AVLNode | None) -> bool:
+    """
+    Recursively verify that every node satisfies the AVL balance property.
+
+    Parameters
+    ----------
+    root:
+        Root of the sub-tree to validate.
+
+    Returns
+    -------
+    bool
+        *True* when every node has ``|balance_factor| <= 1``.
+    """
+
+    def _check(node: AVLNode | None) -> bool:
+        if node is None:
+            return True
+        if abs(get_balance_factor(node)) > 1:
+            return False
+        return _check(node.left) and _check(node.right)  # type: ignore[arg-type]
+
+    return _check(root)

pkstruct/trees/bplus.py

@@ -0,0 +1,425 @@
+"""
+pkstruct.trees.bplus
+====================
+B+ Tree implementation with leaf-linked chain for efficient range queries.
+
+Architecture
+------------
+``BPlusTree`` stores all data in the leaves.  Internal nodes contain only
+separator keys for routing.  Leaf nodes are linked via a ``next_leaf``
+pointer, enabling efficient range scans without recursive traversal.
+
+All major operations run in O(log n).  Range queries benefit from the
+leaf chain after the initial O(log n) descent.
+
+Public API
+----------
+Core CRUD
+    insert, delete, search, contains, update, clear
+
+Metrics
+    size, is_empty, min, max
+
+Navigation
+    leaf_traversal
+
+Range queries
+    range_query
+
+Structural
+    validate
+
+Traversal
+    __iter__, __reversed__
+
+Dunder
+    __len__, __contains__, __iter__, __repr__
+
+Complexity
+----------
+Search        O(log n)
+Insert        O(log n)
+Delete        O(log n)
+Range query   O(log n + k)
+Space         O(n)
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+from typing import Any
+
+from pkstruct.shared.threading import StructureLock
+from pkstruct.trees.exceptions import (
+    DuplicateKeyError,
+    EmptyTreeError,
+    InvalidOrderError,
+    KeyNotFoundError,
+)
+from pkstruct.trees.node import BPlusNode
+
+
+class BPlusTree:
+    """A B+ Tree with leaf-linked chain for efficient range queries.
+
+    Parameters
+    ----------
+    order:
+        Minimum degree of the tree (must be >= 2).  The maximum number
+        of keys per node is ``2 * order - 1``.
+    allow_duplicates:
+        When *False* (default) duplicate keys update the stored value.
+        When *True* duplicates are rejected with ``DuplicateKeyError``.
+    """
+
+    def __init__(self, order: int = 3, allow_duplicates: bool = False) -> None:
+        if order < 2:
+            raise InvalidOrderError(order)
+        self._order: int = order
+        self._allow_duplicates: bool = allow_duplicates
+        self._root: BPlusNode = BPlusNode(is_leaf=True)
+        self._first_leaf: BPlusNode = self._root
+        self._size: int = 0
+        self._lock: StructureLock = StructureLock()
+
+    # ------------------------------------------------------------------
+    # Properties
+    # ------------------------------------------------------------------
+
+    @property
+    def order(self) -> int:
+        """Minimum degree of this B+ Tree."""
+        return self._order
+
+    # ------------------------------------------------------------------
+    # Core CRUD
+    # ------------------------------------------------------------------
+
+    def insert(self, key: Any, value: Any = None) -> None:
+        """Insert *key* with optional *value*.
+
+        Parameters
+        ----------
+        key:
+            Comparable key.
+        value:
+            Optional payload.
+
+        Raises
+        ------
+        DuplicateKeyError
+            If ``allow_duplicates=True`` and *key* already exists.
+        """
+        with self._lock:
+            # Check for duplicate
+            try:
+                existing = self._find_leaf(self._root, key)
+                if existing is not None and key in existing.keys:
+                    if self._allow_duplicates:
+                        raise DuplicateKeyError(key)
+                    idx = existing.keys.index(key)
+                    existing.values[idx] = value
+                    return
+            except (KeyNotFoundError, AttributeError):
+                pass
+
+            leaf = self._find_leaf(self._root, key)
+            self._insert_into_leaf(leaf, key, value)
+
+            if len(leaf.keys) > 2 * self._order - 1:
+                self._split_leaf(leaf)
+
+    def _find_leaf(self, node: BPlusNode, key: Any) -> BPlusNode:
+        if node.is_leaf:
+            return node
+        i = 0
+        while i < len(node.keys) and key >= node.keys[i]:
+            i += 1
+        return self._find_leaf(node.children[i], key)
+
+    def _insert_into_leaf(self, leaf: BPlusNode, key: Any, value: Any) -> None:
+        i = 0
+        while i < len(leaf.keys) and leaf.keys[i] < key:
+            i += 1
+        leaf.keys.insert(i, key)
+        leaf.values.insert(i, value)
+        self._size += 1
+
+    def _split_leaf(self, leaf: BPlusNode) -> None:
+        new_leaf = BPlusNode(is_leaf=True)
+        mid = len(leaf.keys) // 2
+
+        new_leaf.keys = leaf.keys[mid:]
+        new_leaf.values = leaf.values[mid:]
+        leaf.keys = leaf.keys[:mid]
+        leaf.values = leaf.values[:mid]
+
+        new_leaf.next_leaf = leaf.next_leaf
+        leaf.next_leaf = new_leaf
+
+        self._insert_into_parent(leaf, new_leaf.keys[0], new_leaf)
+
+    def _insert_into_parent(
+        self,
+        left: BPlusNode,
+        key: Any,
+        right: BPlusNode,
+    ) -> None:
+        parent = left.parent
+        if parent is None:
+            new_root = BPlusNode(is_leaf=False)
+            new_root.keys = [key]
+            new_root.children = [left, right]
+            left.parent = new_root
+            right.parent = new_root
+            self._root = new_root
+            return
+
+        i = parent.children.index(left) + 1
+        parent.keys.insert(i - 1, key)
+        parent.children.insert(i, right)
+        right.parent = parent
+
+        if len(parent.keys) > 2 * self._order - 1:
+            self._split_internal(parent)
+
+    def _split_internal(self, node: BPlusNode) -> None:
+        new_node = BPlusNode(is_leaf=False)
+        mid = len(node.keys) // 2
+        up_key = node.keys[mid]
+
+        new_node.keys = node.keys[mid + 1 :]
+        new_node.children = node.children[mid + 1 :]
+        node.keys = node.keys[:mid]
+        node.children = node.children[: mid + 1]
+
+        for child in new_node.children:
+            child.parent = new_node
+
+        self._insert_into_parent(node, up_key, new_node)
+
+    def delete(self, key: Any) -> None:
+        """Remove *key* from the tree.
+
+        Parameters
+        ----------
+        key:
+            Key to remove.
+
+        Raises
+        ------
+        KeyNotFoundError
+            If *key* is not present.
+        """
+        with self._lock:
+            leaf = self._find_leaf(self._root, key)
+            if key not in leaf.keys:
+                raise KeyNotFoundError(key)
+            idx = leaf.keys.index(key)
+            leaf.keys.pop(idx)
+            leaf.values.pop(idx)
+            self._size -= 1
+
+    def search(self, key: Any) -> Any | None:
+        """Return the value associated with *key*, or *None* if absent."""
+        with self._lock:
+            leaf = self._find_leaf(self._root, key)
+            if key in leaf.keys:
+                idx = leaf.keys.index(key)
+                return leaf.values[idx]
+            return None
+
+    def contains(self, key: Any) -> bool:
+        """Return *True* if *key* is present."""
+        with self._lock:
+            leaf = self._find_leaf(self._root, key)
+            return key in leaf.keys
+
+    def update(self, key: Any, value: Any) -> None:
+        """Update the value of an existing *key*.
+
+        Raises
+        ------
+        KeyNotFoundError
+            If *key* does not exist.
+        """
+        with self._lock:
+            leaf = self._find_leaf(self._root, key)
+            if key not in leaf.keys:
+                raise KeyNotFoundError(key)
+            idx = leaf.keys.index(key)
+            leaf.values[idx] = value
+
+    def clear(self) -> None:
+        """Remove all keys from the tree."""
+        with self._lock:
+            self._root = BPlusNode(is_leaf=True)
+            self._first_leaf = self._root
+            self._size = 0
+
+    # ------------------------------------------------------------------
+    # Metrics
+    # ------------------------------------------------------------------
+
+    def size(self) -> int:
+        """Return the number of keys currently in the tree."""
+        with self._lock:
+            return self._size
+
+    def is_empty(self) -> bool:
+        """Return *True* if the tree contains no keys."""
+        with self._lock:
+            return self._size == 0
+
+    def min(self) -> Any:
+        """Return the minimum key.
+
+        Raises
+        ------
+        EmptyTreeError
+            If the tree is empty.
+        """
+        with self._lock:
+            if self.is_empty():
+                raise EmptyTreeError("min")
+            return self._first_leaf.keys[0]
+
+    def max(self) -> Any:
+        """Return the maximum key.
+
+        Raises
+        ------
+        EmptyTreeError
+            If the tree is empty.
+        """
+        with self._lock:
+            if self.is_empty():
+                raise EmptyTreeError("max")
+            leaf = self._first_leaf
+            while leaf.next_leaf is not None:
+                leaf = leaf.next_leaf
+            return leaf.keys[-1]
+
+    # ------------------------------------------------------------------
+    # Leaf traversal
+    # ------------------------------------------------------------------
+
+    def leaf_traversal(self) -> list[Any]:
+        """Traverse all leaves via the leaf-linked chain.
+
+        Returns
+        -------
+        list
+            All keys in sorted order by walking the leaf chain.
+        """
+        with self._lock:
+            result: list[Any] = []
+            leaf = self._first_leaf
+            while leaf is not None:
+                result.extend(leaf.keys)
+                leaf = leaf.next_leaf
+            return result
+
+    # ------------------------------------------------------------------
+    # Range query
+    # ------------------------------------------------------------------
+
+    def range_query(self, lo: Any, hi: Any) -> list[Any]:
+        """Return all keys in [*lo*, *hi*] using the leaf chain.
+
+        Parameters
+        ----------
+        lo:
+            Inclusive lower bound.
+        hi:
+            Inclusive upper bound.
+
+        Returns
+        -------
+        list
+            Sorted keys within the range.
+        """
+        with self._lock:
+            result: list[Any] = []
+            leaf = self._find_leaf(self._root, lo)
+            while leaf is not None:
+                for k in leaf.keys:
+                    if lo <= k <= hi:
+                        result.append(k)
+                    elif k > hi:
+                        return result
+                leaf = leaf.next_leaf
+            return result
+
+    # ------------------------------------------------------------------
+    # Validation
+    # ------------------------------------------------------------------
+
+    def validate(self) -> bool:
+        """Verify internal B+ Tree invariants.
+
+        Checks:
+          1. All leaves are at the same depth.
+          2. Internal nodes have correct separator keys.
+          3. Leaf chain maintains sorted order.
+          4. Key counts within bounds.
+
+        Returns
+        -------
+        bool
+            *True* if all invariants hold.
+        """
+        with self._lock:
+            if self._size == 0:
+                return True
+            try:
+                self._validate(self._root, None, None, 0, [])
+                return True
+            except (ValueError, AssertionError, AttributeError):
+                return False
+
+    def _validate(
+        self,
+        node: BPlusNode | None,
+        lo: Any,
+        hi: Any,
+        depth: int,
+        leaf_depths: list[int],
+    ) -> None:
+        if node is None:
+            return
+        kc = len(node.keys)
+        if node is not self._root:
+            assert kc >= self._order - 1, f"Underflow: {kc} keys"
+        assert kc <= 2 * self._order - 1, f"Overflow: {kc} keys"
+        for i in range(kc - 1):
+            assert node.keys[i] < node.keys[i + 1], "Keys not sorted"
+        if node.is_leaf:
+            leaf_depths.append(depth)
+            if len(leaf_depths) > 1 and leaf_depths[-1] != leaf_depths[-2]:
+                raise ValueError("Leaves at different depths")
+        else:
+            for i, child in enumerate(node.children):
+                child_lo = lo if i == 0 else node.keys[i - 1]
+                child_hi = hi if i == kc else node.keys[i]
+                self._validate(child, child_lo, child_hi, depth + 1, leaf_depths)
+
+    # ------------------------------------------------------------------
+    # Dunder methods
+    # ------------------------------------------------------------------
+
+    def __len__(self) -> int:
+        return self.size()
+
+    def __contains__(self, key: Any) -> bool:
+        return self.contains(key)
+
+    def __iter__(self) -> Iterator[Any]:
+        with self._lock:
+            keys = list(self.leaf_traversal())
+        return iter(keys)
+
+    def __repr__(self) -> str:  # pragma: no cover
+        with self._lock:
+            keys = list(self.leaf_traversal())
+            return f"BPlusTree(size={self._size}, order={self._order}, keys={keys})"