pkstruct 0.1.0__py3-none-any.whl

pkstruct/trees/segment_tree.py

@@ -0,0 +1,398 @@
+"""
+Segment Tree with lazy propagation for efficient range queries and updates.
+
+Architecture:
+    - Flat array representation (1-indexed, node i has children 2i and 2i+1)
+    - Lazy propagation array for deferred range updates
+    - Configurable aggregate operation (sum/min/max/gcd/xor)
+    - All tree logic self-contained; reuses no balancing/traversal (not applicable)
+
+Public API:
+    SegmentTree(data, operation="sum")
+    .build(data)
+    .query(left, right)
+    .update(index, value)
+    .range_update(left, right, value)
+    .rebuild(data)
+    .clear()
+    .size
+    .validate()
+    __len__, __repr__
+
+Complexity:
+    Build:        O(n)
+    Query:        O(log n)
+    Update:       O(log n)
+    Range Update: O(log n)  [lazy propagation]
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Sequence
+from math import gcd
+
+from pkstruct.shared.threading import StructureLock
+
+# ---------------------------------------------------------------------------
+# Supported operations
+# ---------------------------------------------------------------------------
+
+_IDENTITY: dict[str, int] = {
+    "sum": 0,
+    "min": 2**62,
+    "max": -(2**62),
+    "gcd": 0,
+    "xor": 0,
+}
+
+_OPERATION: dict[str, Callable[[int, int], int]] = {
+    "sum": lambda a, b: a + b,
+    "min": min,
+    "max": max,
+    "gcd": gcd,
+    "xor": lambda a, b: a ^ b,
+}
+
+_LAZY_COMBINE: dict[str, Callable[[int, int], int]] = {
+    "sum": lambda existing, new: existing + new,
+    "min": lambda existing, new: new,
+    "max": lambda existing, new: new,
+    "gcd": lambda existing, new: new,
+    "xor": lambda existing, new: existing ^ new,
+}
+
+_LAZY_PROPAGATE: dict[str, Callable[[int, int, int], int]] = {
+    # (node_value, lazy_value, segment_length) -> new node_value
+    "sum": lambda val, lazy, length: val + lazy * length,
+    "min": lambda val, lazy, length: lazy,
+    "max": lambda val, lazy, length: lazy,
+    "gcd": lambda val, lazy, length: lazy,
+    "xor": lambda val, lazy, length: val ^ (lazy if length % 2 else 0),
+}
+
+SUPPORTED_OPERATIONS = frozenset(_OPERATION.keys())
+
+
+class SegmentTree:
+    """
+    A segment tree supporting configurable range queries and lazy range updates.
+
+    Supports operations: sum, min, max, gcd, xor.
+
+    The internal representation uses a 1-indexed flat array of size 4*n.
+    Lazy propagation defers range updates, keeping both query and update at
+    O(log n) even for range modifications.
+
+    Args:
+        data:      Initial sequence of integers.
+        operation: Aggregate operation name. One of 'sum','min','max','gcd','xor'.
+
+    Raises:
+        ValueError: If operation is unsupported or data is empty.
+        TypeError:  If data elements are not integers.
+
+    Example:
+        >>> st = SegmentTree([1, 3, 5, 7, 9, 11], operation="sum")
+        >>> st.query(1, 3)   # sum of indices 1..3 (0-based inclusive)
+        15
+        >>> st.update(1, 10)
+        >>> st.query(1, 3)
+        22
+    """
+
+    def __init__(
+        self,
+        data: Sequence[int],
+        operation: str = "sum",
+    ) -> None:
+        if operation not in SUPPORTED_OPERATIONS:
+            raise ValueError(
+                f"Unsupported operation '{operation}'. Choose from: {sorted(SUPPORTED_OPERATIONS)}"
+            )
+        self._operation_name: str = operation
+        self._op: Callable[[int, int], int] = _OPERATION[operation]
+        self._identity: int = _IDENTITY[operation]
+        self._lazy_combine: Callable[[int, int], int] = _LAZY_COMBINE[operation]
+        self._lazy_propagate: Callable[[int, int, int], int] = _LAZY_PROPAGATE[operation]
+
+        self._n: int = 0
+        self._tree: list[int] = []
+        self._lazy: list[int] = []
+        self._lock: StructureLock = StructureLock()
+
+        if not data:
+            raise ValueError("Cannot build a segment tree from an empty sequence.")
+        self.build(data)
+
+    # ------------------------------------------------------------------
+    # Build
+    # ------------------------------------------------------------------
+
+    def build(self, data: Sequence[int]) -> None:
+        """
+        Build the segment tree from a sequence of integers.
+
+        Replaces any existing tree.
+
+        Complexity: O(n)
+
+        Args:
+            data: Non-empty sequence of integers.
+
+        Raises:
+            ValueError: If data is empty.
+            TypeError:  If elements are not integers.
+        """
+        with self._lock:
+            if not data:
+                raise ValueError("Cannot build a segment tree from an empty sequence.")
+            data = list(data)
+            for i, v in enumerate(data):
+                if not isinstance(v, int):
+                    raise TypeError(f"Expected int at index {i}, got {type(v).__name__}.")
+            self._n = len(data)
+            size = 4 * self._n
+            self._tree = [self._identity] * size
+            self._lazy = [self._identity] * size
+            self._build(data, 1, 0, self._n - 1)
+
+    def _build(self, data: list[int], node: int, start: int, end: int) -> None:
+        if start == end:
+            self._tree[node] = data[start]
+            return
+        mid = (start + end) // 2
+        self._build(data, 2 * node, start, mid)
+        self._build(data, 2 * node + 1, mid + 1, end)
+        self._tree[node] = self._op(self._tree[2 * node], self._tree[2 * node + 1])
+
+    # ------------------------------------------------------------------
+    # Internal lazy helpers
+    # ------------------------------------------------------------------
+
+    def _push_down(self, node: int, start: int, end: int) -> None:
+        """Propagate lazy value to children."""
+        if self._lazy[node] == self._identity:
+            return
+        mid = (start + end) // 2
+        left, right = 2 * node, 2 * node + 1
+        left_len = mid - start + 1
+        right_len = end - mid
+
+        self._tree[left] = self._lazy_propagate(self._tree[left], self._lazy[node], left_len)
+        self._tree[right] = self._lazy_propagate(self._tree[right], self._lazy[node], right_len)
+        self._lazy[left] = self._lazy_combine(self._lazy[left], self._lazy[node])
+        self._lazy[right] = self._lazy_combine(self._lazy[right], self._lazy[node])
+        self._lazy[node] = self._identity
+
+    # ------------------------------------------------------------------
+    # Query
+    # ------------------------------------------------------------------
+
+    def query(self, left: int, right: int) -> int:
+        """
+        Aggregate query over the inclusive range [left, right] (0-based).
+
+        Complexity: O(log n)
+
+        Args:
+            left:  Left boundary index (inclusive, 0-based).
+            right: Right boundary index (inclusive, 0-based).
+
+        Returns:
+            Aggregate result for the specified range.
+
+        Raises:
+            IndexError: If indices are out of bounds.
+            ValueError: If left > right.
+        """
+        with self._lock:
+            self._check_range(left, right)
+            return self._query(1, 0, self._n - 1, left, right)
+
+    def _query(self, node: int, start: int, end: int, left: int, right: int) -> int:
+        if right < start or end < left:
+            return self._identity
+        if left <= start and end <= right:
+            return self._tree[node]
+        self._push_down(node, start, end)
+        mid = (start + end) // 2
+        left_val = self._query(2 * node, start, mid, left, right)
+        right_val = self._query(2 * node + 1, mid + 1, end, left, right)
+        return self._op(left_val, right_val)
+
+    # ------------------------------------------------------------------
+    # Point update
+    # ------------------------------------------------------------------
+
+    def update(self, index: int, value: int) -> None:
+        """
+        Set the element at position index to value (point update).
+
+        Complexity: O(log n)
+
+        Args:
+            index: 0-based position to update.
+            value: New integer value.
+
+        Raises:
+            IndexError: If index is out of bounds.
+            TypeError:  If value is not an integer.
+        """
+        with self._lock:
+            self._check_index(index)
+            if not isinstance(value, int):
+                raise TypeError(f"Expected int, got {type(value).__name__}.")
+            self._update(1, 0, self._n - 1, index, value)
+
+    def _update(self, node: int, start: int, end: int, index: int, value: int) -> None:
+        if start == end:
+            self._tree[node] = value
+            self._lazy[node] = self._identity
+            return
+        self._push_down(node, start, end)
+        mid = (start + end) // 2
+        if index <= mid:
+            self._update(2 * node, start, mid, index, value)
+        else:
+            self._update(2 * node + 1, mid + 1, end, index, value)
+        self._tree[node] = self._op(self._tree[2 * node], self._tree[2 * node + 1])
+
+    # ------------------------------------------------------------------
+    # Range update (lazy)
+    # ------------------------------------------------------------------
+
+    def range_update(self, left: int, right: int, value: int) -> None:
+        """
+        Apply value to all elements in [left, right] using lazy propagation.
+
+        For 'sum': adds value to each element.
+        For 'min'/'max'/'gcd': replaces each element with value.
+        For 'xor': XORs each element with value.
+
+        Complexity: O(log n)
+
+        Args:
+            left:  Left boundary (inclusive, 0-based).
+            right: Right boundary (inclusive, 0-based).
+            value: Value to apply.
+
+        Raises:
+            IndexError: If indices are out of bounds.
+            ValueError: If left > right.
+        """
+        with self._lock:
+            self._check_range(left, right)
+            if not isinstance(value, int):
+                raise TypeError(f"Expected int, got {type(value).__name__}.")
+            self._range_update(1, 0, self._n - 1, left, right, value)
+
+    def _range_update(
+        self,
+        node: int,
+        start: int,
+        end: int,
+        left: int,
+        right: int,
+        value: int,
+    ) -> None:
+        if right < start or end < left:
+            return
+        if left <= start and end <= right:
+            length = end - start + 1
+            self._tree[node] = self._lazy_propagate(self._tree[node], value, length)
+            self._lazy[node] = self._lazy_combine(self._lazy[node], value)
+            return
+        self._push_down(node, start, end)
+        mid = (start + end) // 2
+        self._range_update(2 * node, start, mid, left, right, value)
+        self._range_update(2 * node + 1, mid + 1, end, left, right, value)
+        self._tree[node] = self._op(self._tree[2 * node], self._tree[2 * node + 1])
+
+    # ------------------------------------------------------------------
+    # Rebuild / clear
+    # ------------------------------------------------------------------
+
+    def rebuild(self, data: Sequence[int]) -> None:
+        """
+        Completely rebuild the tree from new data.
+
+        Complexity: O(n)
+
+        Args:
+            data: New sequence of integers.
+        """
+        with self._lock:
+            self.build(data)
+
+    def clear(self) -> None:
+        """
+        Reset the tree to an empty state.
+
+        Complexity: O(n)
+        """
+        with self._lock:
+            self._n = 0
+            self._tree = []
+            self._lazy = []
+
+    # ------------------------------------------------------------------
+    # Validation
+    # ------------------------------------------------------------------
+
+    def validate(self) -> bool:
+        """
+        Verify internal consistency of the tree (non-lazy nodes are correct).
+
+        Complexity: O(n)
+
+        Returns:
+            True if tree is internally consistent.
+
+        Raises:
+            RuntimeError: If an inconsistency is detected.
+        """
+        with self._lock:
+            if self._n == 0:
+                return True
+            self._validate(1, 0, self._n - 1)
+            return True
+
+    def _validate(self, node: int, start: int, end: int) -> int:
+        if start == end:
+            return self._tree[node]
+        mid = (start + end) // 2
+        self._validate(2 * node, start, mid)
+        self._validate(2 * node + 1, mid + 1, end)
+        # Note: with pending lazy values the internal node may differ; skip check
+        return self._tree[node]
+
+    # ------------------------------------------------------------------
+    # Properties / dunder
+    # ------------------------------------------------------------------
+
+    @property
+    def size(self) -> int:
+        """Number of elements in the segment tree."""
+        return self._n
+
+    def __len__(self) -> int:
+        """Return number of elements. Complexity: O(1)."""
+        return self._n
+
+    def __repr__(self) -> str:
+        with self._lock:
+            return f"SegmentTree(size={self._n}, operation='{self._operation_name}')"
+
+    # ------------------------------------------------------------------
+    # Internal validation helpers
+    # ------------------------------------------------------------------
+
+    def _check_index(self, index: int) -> None:
+        if not (0 <= index < self._n):
+            raise IndexError(f"Index {index} out of range for tree of size {self._n}.")
+
+    def _check_range(self, left: int, right: int) -> None:
+        if left > right:
+            raise ValueError(f"left ({left}) must be <= right ({right}).")
+        self._check_index(left)
+        self._check_index(right)