1
1
from __future__ import annotations
2
2
3
+
3
4
class Node :
4
5
def __init__ (self , value : int = 0 ) -> None :
5
6
self .value : int = value
6
7
self .left : Node | None = None
7
8
self .right : Node | None = None
8
9
10
+
9
11
class PersistentSegmentTree :
10
12
def __init__ (self , arr : list [int ]) -> None :
11
13
self .n : int = len (arr )
@@ -81,7 +83,9 @@ def _update(self, node: Node, start: int, end: int, index: int, value: int) -> N
81
83
new_node .left = node .left # Ensure left node is the same as the original
82
84
new_node .right = self ._update (node .right , mid + 1 , end , index , value )
83
85
84
- new_node .value = new_node .left .value + (new_node .right .value if new_node .right else 0 )
86
+ new_node .value = new_node .left .value + (
87
+ new_node .right .value if new_node .right else 0
88
+ )
85
89
86
90
return new_node
87
91
@@ -97,7 +101,7 @@ def query(self, version: int, left: int, right: int) -> int:
97
101
>>> version_1 = pst.update(0, 1, 5) # Update index 1 to 5
98
102
>>> pst.query(version_1, 0, 3) # Sum of all elements in new version
99
103
13
100
- >>> pst.query(version_1, 1, 2) # Sum of elements at index 1 and 2
104
+ >>> pst.query(version_1, 1, 2) # Sum of elements at index 1 and 2
101
105
8
102
106
"""
103
107
return self ._query (self .roots [version ], 0 , self .n - 1 , left , right )
@@ -120,12 +124,15 @@ def _query(self, node: Node, start: int, end: int, left: int, right: int) -> int
120
124
if left <= start and right >= end :
121
125
return node .value
122
126
mid = (start + end ) // 2
123
- return (self ._query (node .left , start , mid , left , right ) +
124
- self ._query (node .right , mid + 1 , end , left , right ))
127
+ return self ._query (node .left , start , mid , left , right ) + self ._query (
128
+ node .right , mid + 1 , end , left , right
129
+ )
130
+
125
131
126
132
# Running the doctests
127
133
if __name__ == "__main__" :
128
134
import doctest
135
+
129
136
print ("Running doctests..." )
130
137
result = doctest .testmod ()
131
138
print (f"Ran { result .attempted } tests, { result .failed } failed." )
0 commit comments