From a187c993ea5295c8333fd1dd92cab4b075838d9f Mon Sep 17 00:00:00 2001 From: vashprog Date: Tue, 7 Oct 2025 21:17:16 +0530 Subject: [PATCH 1/2] Add Segment Tree implementation in Python --- data_structures/trie/segment_tree.py | 100 +++++++++++++++++++++++++++ 1 file changed, 100 insertions(+) create mode 100644 data_structures/trie/segment_tree.py diff --git a/data_structures/trie/segment_tree.py b/data_structures/trie/segment_tree.py new file mode 100644 index 000000000000..40c180f73927 --- /dev/null +++ b/data_structures/trie/segment_tree.py @@ -0,0 +1,100 @@ +""" +A Segment Tree is a binary tree data structure used for efficiently answering +range queries and updates on an array, such as sum, minimum, or maximum over +a subrange. It offers O(log n) time complexity for both queries and updates, +making it very efficient compared to a naive O(n) approach. + +While building the tree takes O(n) time and the tree requires O(n) space, +this preprocessing enables fast range queries that would otherwise be slow. +Segment Trees are especially useful when the array is mutable and queries +and updates are intermixed. + +Time Complexity: +- Build: O(n) +- Query: O(log n) +- Update: O(log n) + +Example usage and doctests: + +>>> data = [1, 2, 3, 4, 5] +>>> st = SegmentTree(data) +>>> st.query(1, 4) +9 +>>> st.update(2, 10) +>>> st.query(1, 4) +16 +""" + + +class SegmentTree: + """Segment Tree for efficient range sum queries.""" + + def __init__(self, data: list[int]): + """Initialize the segment tree with the input data. + + Args: + data (list[int]): List of integers to build the segment tree. + """ + self.n = len(data) + self.tree = [0] * (2 * self.n) + # Build the tree + for i in range(self.n): + self.tree[self.n + i] = data[i] + for i in range(self.n - 1, 0, -1): + self.tree[i] = self.tree[i << 1] + self.tree[i << 1 | 1] + + def update(self, index: int, value: int) -> None: + """Update element at index with a new value. + + Args: + index (int): Index of the element to update. + value (int): New value to set at the given index. + """ + if index < 0 or index >= self.n: + raise ValueError("Index out of bounds") + index += self.n + self.tree[index] = value + while index > 1: + index >>= 1 + self.tree[index] = self.tree[index << 1] + self.tree[index << 1 | 1] + + def query(self, left: int, right: int) -> int: + """Compute the sum of elements in the interval [left, right). + + Args: + left (int): Left index (inclusive). + right (int): Right index (exclusive). + + Returns: + int: Sum of elements from left to right-1. + + Raises: + ValueError: If indices are out of bounds or left >= right. + """ + if left < 0 or right > self.n or left >= right: + raise ValueError("Invalid query range") + res = 0 + left += self.n + right += self.n + while left < right: + if left & 1: + res += self.tree[left] + left += 1 + if right & 1: + right -= 1 + res += self.tree[right] + left >>= 1 + right >>= 1 + return res + + +if __name__ == "__main__": + import doctest + + data = [1, 2, 3, 4, 5] + st = SegmentTree(data) + print("Initial sum 1-4:", st.query(1, 4)) + st.update(2, 10) + print("Updated sum 1-4:", st.query(1, 4)) + + doctest.testmod() From 92576748d9844aa421ca1d6ccb5427a07f237bb3 Mon Sep 17 00:00:00 2001 From: vashprog Date: Wed, 8 Oct 2025 19:53:44 +0530 Subject: [PATCH 2/2] Add Segment Tree implementation for range sum queries with type hints and doctests --- data_structures/trie/segment_tree.py | 157 ++++++++++++++------------- 1 file changed, 82 insertions(+), 75 deletions(-) diff --git a/data_structures/trie/segment_tree.py b/data_structures/trie/segment_tree.py index 40c180f73927..5a476aa4eb32 100644 --- a/data_structures/trie/segment_tree.py +++ b/data_structures/trie/segment_tree.py @@ -1,100 +1,107 @@ """ -A Segment Tree is a binary tree data structure used for efficiently answering -range queries and updates on an array, such as sum, minimum, or maximum over -a subrange. It offers O(log n) time complexity for both queries and updates, -making it very efficient compared to a naive O(n) approach. - -While building the tree takes O(n) time and the tree requires O(n) space, -this preprocessing enables fast range queries that would otherwise be slow. -Segment Trees are especially useful when the array is mutable and queries -and updates are intermixed. +A Segment Tree is a binary tree used for storing intervals or segments. +It allows querying the sum, minimum, or maximum of elements in a range efficiently. Time Complexity: - Build: O(n) - Query: O(log n) - Update: O(log n) - -Example usage and doctests: - ->>> data = [1, 2, 3, 4, 5] ->>> st = SegmentTree(data) ->>> st.query(1, 4) -9 ->>> st.update(2, 10) ->>> st.query(1, 4) -16 """ class SegmentTree: - """Segment Tree for efficient range sum queries.""" - - def __init__(self, data: list[int]): - """Initialize the segment tree with the input data. - - Args: - data (list[int]): List of integers to build the segment tree. + def __init__(self, data: list[int]) -> None: + """ + Initializes the Segment Tree from a list of integers. + :param data: List of integer elements. + :return: None """ self.n = len(data) - self.tree = [0] * (2 * self.n) - # Build the tree - for i in range(self.n): - self.tree[self.n + i] = data[i] - for i in range(self.n - 1, 0, -1): - self.tree[i] = self.tree[i << 1] + self.tree[i << 1 | 1] + self.tree = [0] * (4 * self.n) + self._build(data, 0, 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] + else: + mid = (start + end) // 2 + self._build(data, 2 * node + 1, start, mid) + self._build(data, 2 * node + 2, mid + 1, end) + self.tree[node] = self.tree[2 * node + 1] + self.tree[2 * node + 2] def update(self, index: int, value: int) -> None: - """Update element at index with a new value. - - Args: - index (int): Index of the element to update. - value (int): New value to set at the given index. """ - if index < 0 or index >= self.n: - raise ValueError("Index out of bounds") - index += self.n - self.tree[index] = value - while index > 1: - index >>= 1 - self.tree[index] = self.tree[index << 1] + self.tree[index << 1 | 1] + Updates the value at the given index and updates the tree. + :param index: Index to update. + :param value: New value. + :return: None + """ + self._update(0, 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 + else: + mid = (start + end) // 2 + if index <= mid: + self._update(2 * node + 1, start, mid, index, value) + else: + self._update(2 * node + 2, mid + 1, end, index, value) + self.tree[node] = self.tree[2 * node + 1] + self.tree[2 * node + 2] def query(self, left: int, right: int) -> int: - """Compute the sum of elements in the interval [left, right). - - Args: - left (int): Left index (inclusive). - right (int): Right index (exclusive). - - Returns: - int: Sum of elements from left to right-1. - - Raises: - ValueError: If indices are out of bounds or left >= right. """ - if left < 0 or right > self.n or left >= right: - raise ValueError("Invalid query range") - res = 0 - left += self.n - right += self.n - while left < right: - if left & 1: - res += self.tree[left] - left += 1 - if right & 1: - right -= 1 - res += self.tree[right] - left >>= 1 - right >>= 1 - return res + Returns the sum of elements in the range [left, right]. + :param left: Left index (inclusive). + :param right: Right index (inclusive). + :return: Sum of the range. + + >>> data = [1, 2, 3, 4, 5] + >>> st = SegmentTree(data) + >>> st.query(1, 3) + 9 + >>> st.update(2, 10) + >>> st.query(1, 3) + 16 + """ + return self._query(0, 0, self.n - 1, left, right) + def _query(self, node: int, start: int, end: int, left: int, right: int) -> int: + if right < start or left > end: + return 0 + if left <= start and end <= right: + return self.tree[node] + mid = (start + end) // 2 + return self._query(2 * node + 1, start, mid, left, right) + self._query( + 2 * node + 2, mid + 1, end, left, right + ) -if __name__ == "__main__": - import doctest +def test_segment_tree() -> bool: data = [1, 2, 3, 4, 5] st = SegmentTree(data) - print("Initial sum 1-4:", st.query(1, 4)) + assert st.query(0, 2) == 6 + assert st.query(1, 4) == 14 st.update(2, 10) - print("Updated sum 1-4:", st.query(1, 4)) + assert st.query(1, 3) == 16 + assert st.query(0, 4) == 22 + return True + + +def print_results(msg: str, passes: bool) -> None: + print(str(msg), "works!" if passes else "doesn't work :(") + - doctest.testmod() +def pytests() -> None: + assert test_segment_tree() + + +def main() -> None: + """ + >>> pytests() + """ + print_results("Testing Segment Tree functionality", test_segment_tree()) + + +if __name__ == "__main__": + main()