33
44class SegmentTree :
55 def __init__ (self , a ):
6- self .N = len (a )
6+ self .A = a
7+ self .N = len (self .A )
78 self .st = [0 ] * (
89 4 * self .N
910 ) # approximate the overall size of segment tree with array N
1011 if self .N :
1112 self .build (1 , 0 , self .N - 1 )
1213
1314 def left (self , idx ):
15+ """
16+ Returns the left child index for a given index in a binary tree.
17+
18+ >>> s = SegmentTree([1, 2, 3])
19+ >>> s.left(1)
20+ 2
21+ >>> s.left(2)
22+ 4
23+ """
1424 return idx * 2
1525
1626 def right (self , idx ):
27+ """
28+ Returns the right child index for a given index in a binary tree.
29+
30+ >>> s = SegmentTree([1, 2, 3])
31+ >>> s.right(1)
32+ 3
33+ >>> s.right(2)
34+ 5
35+ """
1736 return idx * 2 + 1
1837
1938 def build (self , idx , l , r ): # noqa: E741
2039 if l == r :
21- self .st [idx ] = A [l ]
40+ self .st [idx ] = self . A [l ]
2241 else :
2342 mid = (l + r ) // 2
2443 self .build (self .left (idx ), l , mid )
2544 self .build (self .right (idx ), mid + 1 , r )
2645 self .st [idx ] = max (self .st [self .left (idx )], self .st [self .right (idx )])
2746
2847 def update (self , a , b , val ):
48+ """
49+ Update the values in the segment tree in the range [a,b] with the given value.
50+
51+ >>> s = SegmentTree([1, 2, 3, 4, 5])
52+ >>> s.update(2, 4, 10)
53+ True
54+ >>> s.query(1, 5)
55+ 10
56+ """
2957 return self .update_recursive (1 , 0 , self .N - 1 , a - 1 , b - 1 , val )
3058
3159 def update_recursive (self , idx , l , r , a , b , val ): # noqa: E741
@@ -44,6 +72,15 @@ def update_recursive(self, idx, l, r, a, b, val): # noqa: E741
4472 return True
4573
4674 def query (self , a , b ):
75+ """
76+ Query the maximum value in the range [a,b].
77+
78+ >>> s = SegmentTree([1, 2, 3, 4, 5])
79+ >>> s.query(1, 3)
80+ 3
81+ >>> s.query(1, 5)
82+ 5
83+ """
4784 return self .query_recursive (1 , 0 , self .N - 1 , a - 1 , b - 1 )
4885
4986 def query_recursive (self , idx , l , r , a , b ): # noqa: E741
0 commit comments