diff --git a/data_structures/fenwick_tree/fenwick_tree.py b/data_structures/fenwick_tree/fenwick_tree.py new file mode 100644 index 000000000000..7a2afb4b92ad --- /dev/null +++ b/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() diff --git a/data_structures/segment_tree/segment_tree.py b/data_structures/segment_tree/segment_tree.py new file mode 100644 index 000000000000..b8b5b2996bb9 --- /dev/null +++ b/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()