From 04b44c4c7b78aabf2acad2e59d4a64c405ded94b Mon Sep 17 00:00:00 2001
From: Putul Singh <127419636+putul03@users.noreply.github.com>
Date: Sat, 19 Oct 2024 16:22:59 +0530
Subject: [PATCH] Update persistent_segment_tree.py
---
data_structures/persistent_segment_tree.py | 33 +++++++++++++++++++---
1 file changed, 29 insertions(+), 4 deletions(-)
diff --git a/data_structures/persistent_segment_tree.py b/data_structures/persistent_segment_tree.py
index 32f1ec0df..b0bc7ed3f 100644
--- a/data_structures/persistent_segment_tree.py
+++ b/data_structures/persistent_segment_tree.py
@@ -34,6 +34,19 @@ class PersistentSegmentTree:
return node
def update(self, version: int, index: int, value: int) -> int:
+ """
+ Updates the value at the given index and returns the new version.
+
+ >>> pst = PersistentSegmentTree([1, 2, 3, 4])
+ >>> version_1 = pst.update(0, 1, 5) # Update index 1 to 5
+ >>> pst.query(version_1, 0, 3) # Query sum of all elements in new version
+ 13
+ >>> pst.query(0, 0, 3) # Original version remains unchanged
+ 10
+ >>> version_2 = pst.update(version_1, 3, 6) # Update index 3 to 6 in version_1
+ >>> pst.query(version_2, 0, 3) # Query sum of all elements in newest version
+ 15
+ """
new_root = self._update(self.roots[version], 0, self.n - 1, index, value)
self.roots.append(new_root)
return len(self.roots) - 1
@@ -72,6 +85,20 @@ class PersistentSegmentTree:
return new_node
def query(self, version: int, left: int, right: int) -> int:
+ """
+ Queries the sum in the given range for the specified version.
+
+ >>> pst = PersistentSegmentTree([1, 2, 3, 4])
+ >>> pst.query(0, 0, 3) # Sum of all elements in original version
+ 10
+ >>> pst.query(0, 1, 2) # Sum of elements at index 1 and 2 in original version
+ 5
+ >>> version_1 = pst.update(0, 1, 5) # Update index 1 to 5
+ >>> pst.query(version_1, 0, 3) # Sum of all elements in new version
+ 13
+ >>> pst.query(version_1, 1, 2) # Sum of elements at index 1 and 2 in new version
+ 8
+ """
return self._query(self.roots[version], 0, self.n - 1, left, right)
def _query(self, node: Node, start: int, end: int, left: int, right: int) -> int:
@@ -92,15 +119,13 @@ class PersistentSegmentTree:
if left <= start and right >= end:
return node.value
mid = (start + end) // 2
- return self._query(node.left, start, mid, left, right) + self._query(
- node.right, mid + 1, end, left, right
- )
+ return (self._query(node.left, start, mid, left, right) +
+ self._query(node.right, mid + 1, end, left, right))
# Running the doctests
if __name__ == "__main__":
import doctest
-
print("Running doctests...")
result = doctest.testmod()
print(f"Ran {result.attempted} tests, {result.failed} failed.")