Add Segment Tree data structure implementation

Implement Segment Tree for range queries and updates.

Author: Tejasrahane

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()