3
3
4
4
class SegmentTree :
5
5
def __init__ (self , a ):
6
- self .N = len (a )
6
+ self .A = a
7
+ self .N = len (self .A )
7
8
self .st = [0 ] * (
8
9
4 * self .N
9
10
) # approximate the overall size of segment tree with array N
10
11
if self .N :
11
12
self .build (1 , 0 , self .N - 1 )
12
13
13
14
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
+ """
14
24
return idx * 2
15
25
16
26
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
+ """
17
36
return idx * 2 + 1
18
37
19
38
def build (self , idx , l , r ): # noqa: E741
20
39
if l == r :
21
- self .st [idx ] = A [l ]
40
+ self .st [idx ] = self . A [l ]
22
41
else :
23
42
mid = (l + r ) // 2
24
43
self .build (self .left (idx ), l , mid )
25
44
self .build (self .right (idx ), mid + 1 , r )
26
45
self .st [idx ] = max (self .st [self .left (idx )], self .st [self .right (idx )])
27
46
28
47
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
+ """
29
57
return self .update_recursive (1 , 0 , self .N - 1 , a - 1 , b - 1 , val )
30
58
31
59
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
44
72
return True
45
73
46
74
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
+ """
47
84
return self .query_recursive (1 , 0 , self .N - 1 , a - 1 , b - 1 )
48
85
49
86
def query_recursive (self , idx , l , r , a , b ): # noqa: E741
0 commit comments