From d071d8323bd7ed5aef27a966bf86b869d1beb40a Mon Sep 17 00:00:00 2001 From: Tejasrahane <161036451+Tejasrahane@users.noreply.github.com> Date: Wed, 22 Oct 2025 07:51:08 +0530 Subject: [PATCH 1/2] Add Segment Tree data structure implementation Implement Segment Tree for range queries and updates. --- data_structures/segment_tree/segment_tree.py | 177 +++++++++++++++++++ 1 file changed, 177 insertions(+) create mode 100644 data_structures/segment_tree/segment_tree.py 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() From b72140a5620bd60d73129bdcc2c61f42d9d3b242 Mon Sep 17 00:00:00 2001 From: Tejasrahane <161036451+Tejasrahane@users.noreply.github.com> Date: Wed, 22 Oct 2025 07:52:47 +0530 Subject: [PATCH 2/2] Implement Fenwick Tree with basic operations This implementation supports efficient prefix sum queries and point updates. It includes methods for updating values, calculating prefix sums, range sums, and setting specific values in the Fenwick Tree. --- data_structures/fenwick_tree/fenwick_tree.py | 153 +++++++++++++++++++ 1 file changed, 153 insertions(+) create mode 100644 data_structures/fenwick_tree/fenwick_tree.py 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()