Add data structure programs

data_structures/fenwick_tree/fenwick_tree.py

@@ -0,0 +1,153 @@
+"""Fenwick Tree (Binary Indexed Tree) Data Structure.
+
+A Fenwick Tree is a data structure that can efficiently update elements
+and calculate prefix sums in a table of numbers. It is also known as a
+Binary Indexed Tree (BIT).
+
+Time Complexity:
+- Update: O(log n)
+- Query (prefix sum): O(log n)
+- Construction: O(n log n)
+
+Space Complexity: O(n)
+"""
+
+
+class FenwickTree:
+    """Fenwick Tree implementation for range sum queries.
+
+    This implementation supports efficient prefix sum queries and point updates.
+    The tree uses 1-based indexing internally for simpler implementation.
+
+    Attributes:
+        size: Size of the input array
+        tree: List storing the Fenwick tree values
+
+    >>> ft = FenwickTree(5)
+    >>> ft.update(0, 3)
+    >>> ft.update(1, 2)
+    >>> ft.update(2, 5)
+    >>> ft.update(3, 1)
+    >>> ft.update(4, 4)
+    >>> ft.prefix_sum(2)
+    10
+    >>> ft.range_sum(1, 3)
+    8
+    >>> ft2 = FenwickTree([1, 3, 5, 7, 9])
+    >>> ft2.prefix_sum(2)
+    9
+    >>> ft2.range_sum(1, 3)
+    15
+    """
+
+    def __init__(self, size_or_array: int | list[int]) -> None:
+        """Initialize Fenwick Tree.
+
+        Args:
+            size_or_array: Either size of array (int) or initial array values (list)
+
+        >>> ft = FenwickTree(5)
+        >>> len(ft.tree)
+        6
+        >>> ft2 = FenwickTree([1, 2, 3])
+        >>> ft2.prefix_sum(2)
+        6
+        """
+        if isinstance(size_or_array, int):
+            self.size = size_or_array
+            self.tree = [0] * (self.size + 1)  # 1-indexed
+        else:
+            self.size = len(size_or_array)
+            self.tree = [0] * (self.size + 1)
+            for i, val in enumerate(size_or_array):
+                self.update(i, val)
+
+    def update(self, index: int, delta: int) -> None:
+        """Add delta to element at given index.
+
+        Args:
+            index: Index to update (0-based)
+            delta: Value to add to the element
+
+        >>> ft = FenwickTree(5)
+        >>> ft.update(0, 5)
+        >>> ft.prefix_sum(0)
+        5
+        >>> ft.update(0, 3)
+        >>> ft.prefix_sum(0)
+        8
+        """
+        index += 1  # Convert to 1-based indexing
+        while index <= self.size:
+            self.tree[index] += delta
+            index += index & (-index)  # Add last set bit
+
+    def prefix_sum(self, index: int) -> int:
+        """Calculate sum of elements from 0 to index (inclusive).
+
+        Args:
+            index: End index for prefix sum (0-based, inclusive)
+
+        Returns:
+            Sum of elements from index 0 to given index
+
+        >>> ft = FenwickTree([1, 2, 3, 4, 5])
+        >>> ft.prefix_sum(0)
+        1
+        >>> ft.prefix_sum(2)
+        6
+        >>> ft.prefix_sum(4)
+        15
+        """
+        index += 1  # Convert to 1-based indexing
+        result = 0
+        while index > 0:
+            result += self.tree[index]
+            index -= index & (-index)  # Remove last set bit
+        return result
+
+    def range_sum(self, left: int, right: int) -> int:
+        """Calculate sum of elements in range [left, right].
+
+        Args:
+            left: Left boundary of range (0-based, inclusive)
+            right: Right boundary of range (0-based, inclusive)
+
+        Returns:
+            Sum of elements in the given range
+
+        >>> ft = FenwickTree([1, 2, 3, 4, 5])
+        >>> ft.range_sum(0, 2)
+        6
+        >>> ft.range_sum(2, 4)
+        12
+        >>> ft.range_sum(1, 1)
+        2
+        """
+        if left == 0:
+            return self.prefix_sum(right)
+        return self.prefix_sum(right) - self.prefix_sum(left - 1)
+
+    def set_value(self, index: int, value: int) -> None:
+        """Set element at index to given value.
+
+        Args:
+            index: Index to set (0-based)
+            value: New value to set
+
+        >>> ft = FenwickTree([1, 2, 3, 4, 5])
+        >>> ft.set_value(2, 10)
+        >>> ft.prefix_sum(2)
+        13
+        >>> ft.range_sum(2, 2)
+        10
+        """
+        current_value = self.range_sum(index, index)
+        delta = value - current_value
+        self.update(index, delta)
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod()

data_structures/segment_tree/segment_tree.py

@@ -0,0 +1,177 @@
+"""Segment Tree Data Structure.
+
+A Segment Tree is a binary tree used for storing intervals or segments.
+It allows querying which of the stored segments contain a given point.
+Typically used for range queries and updates.
+
+Time Complexity:
+- Build: O(n)
+- Query: O(log n)
+- Update: O(log n)
+
+Space Complexity: O(n)
+"""
+
+from typing import Callable
+
+
+class SegmentTree:
+    """Segment Tree implementation for range queries.
+
+    This implementation supports range sum queries and point updates.
+    Can be extended to support other operations like min/max queries.
+
+    Attributes:
+        tree: List storing the segment tree nodes
+        n: Size of the input array
+        operation: Function to combine two values (default: addition)
+
+    >>> st = SegmentTree([1, 3, 5, 7, 9, 11])
+    >>> st.query(1, 3)
+    15
+    >>> st.update(1, 10)
+    >>> st.query(1, 3)
+    22
+    >>> st.query(0, 5)
+    42
+    >>> st2 = SegmentTree([2, 4, 6, 8], operation=min)
+    >>> st2.query(0, 3)
+    2
+    >>> st2.update(0, 10)
+    >>> st2.query(0, 3)
+    4
+    """
+
+    def __init__(
+        self, arr: list[int], operation: Callable[[int, int], int] = lambda a, b: a + b
+    ) -> None:
+        """Initialize segment tree with given array.
+
+        Args:
+            arr: Input array of integers
+            operation: Binary operation to combine values (default: addition)
+
+        >>> st = SegmentTree([1, 2, 3])
+        >>> len(st.tree)
+        8
+        """
+        self.n = len(arr)
+        self.tree = [0] * (4 * self.n)  # Allocate space for segment tree
+        self.operation = operation
+        self._build(arr, 0, 0, self.n - 1)
+
+    def _build(self, arr: list[int], node: int, start: int, end: int) -> None:
+        """Build segment tree recursively.
+
+        Args:
+            arr: Input array
+            node: Current node index in tree
+            start: Start index of current segment
+            end: End index of current segment
+        """
+        if start == end:
+            # Leaf node
+            self.tree[node] = arr[start]
+        else:
+            mid = (start + end) // 2
+            left_child = 2 * node + 1
+            right_child = 2 * node + 2
+            self._build(arr, left_child, start, mid)
+            self._build(arr, right_child, mid + 1, end)
+            self.tree[node] = self.operation(
+                self.tree[left_child], self.tree[right_child]
+            )
+
+    def query(self, left: int, right: int) -> int:
+        """Query for value in range [left, right].
+
+        Args:
+            left: Left boundary of query range (inclusive)
+            right: Right boundary of query range (inclusive)
+
+        Returns:
+            Result of applying operation over the range
+
+        >>> st = SegmentTree([1, 2, 3, 4, 5])
+        >>> st.query(0, 2)
+        6
+        >>> st.query(2, 4)
+        12
+        """
+        return self._query(0, 0, self.n - 1, left, right)
+
+    def _query(self, node: int, start: int, end: int, left: int, right: int) -> int:
+        """Recursive helper for range query.
+
+        Args:
+            node: Current node index
+            start: Start of current segment
+            end: End of current segment
+            left: Query left boundary
+            right: Query right boundary
+
+        Returns:
+            Query result for current segment
+        """
+        if right < start or left > end:
+            # No overlap
+            return 0 if self.operation(0, 0) == 0 else float('inf')
+
+        if left <= start and end <= right:
+            # Complete overlap
+            return self.tree[node]
+
+        # Partial overlap
+        mid = (start + end) // 2
+        left_child = 2 * node + 1
+        right_child = 2 * node + 2
+        left_result = self._query(left_child, start, mid, left, right)
+        right_result = self._query(right_child, mid + 1, end, left, right)
+        return self.operation(left_result, right_result)
+
+    def update(self, index: int, value: int) -> None:
+        """Update value at given index.
+
+        Args:
+            index: Index to update
+            value: New value
+
+        >>> st = SegmentTree([1, 2, 3, 4, 5])
+        >>> st.query(0, 4)
+        15
+        >>> st.update(2, 10)
+        >>> st.query(0, 4)
+        22
+        """
+        self._update(0, 0, self.n - 1, index, value)
+
+    def _update(self, node: int, start: int, end: int, index: int, value: int) -> None:
+        """Recursive helper for point update.
+
+        Args:
+            node: Current node index
+            start: Start of current segment
+            end: End of current segment
+            index: Index to update
+            value: New value
+        """
+        if start == end:
+            # Leaf node
+            self.tree[node] = value
+        else:
+            mid = (start + end) // 2
+            left_child = 2 * node + 1
+            right_child = 2 * node + 2
+            if index <= mid:
+                self._update(left_child, start, mid, index, value)
+            else:
+                self._update(right_child, mid + 1, end, index, value)
+            self.tree[node] = self.operation(
+                self.tree[left_child], self.tree[right_child]
+            )
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod()