1
1
import math
2
+ from typing import List
2
3
3
4
4
5
class SegmentTree :
5
- def __init__ (self , N ) :
6
+ def __init__ (self , N : int ) -> None :
6
7
self .N = N
7
- self . st = [
8
- 0 for i in range (0 , 4 * N )
9
- ] # approximate the overall size of segment tree with array N
10
- self .lazy = [0 for i in range (0 , 4 * N )] # create array to store lazy update
11
- self .flag = [0 for i in range (0 , 4 * N )] # flag for lazy update
8
+ # approximate the overall size of segment tree with array N
9
+ self . st : List [ int ] = [ 0 for i in range (0 , 4 * N )]
10
+ # create array to store lazy update
11
+ self .lazy : List [ int ] = [0 for i in range (0 , 4 * N )]
12
+ self .flag : List [ int ] = [0 for i in range (0 , 4 * N )] # flag for lazy update
12
13
13
- def left (self , idx ):
14
+ def left (self , idx : int ) -> int :
15
+ """
16
+ >>> segment_tree = SegmentTree(15)
17
+ >>> segment_tree.left(1)
18
+ 2
19
+ >>> segment_tree.left(2)
20
+ 4
21
+ >>> segment_tree.left(12)
22
+ 24
23
+ """
14
24
return idx * 2
15
25
16
- def right (self , idx ):
26
+ def right (self , idx : int ) -> int :
27
+ """
28
+ >>> segment_tree = SegmentTree(15)
29
+ >>> segment_tree.right(1)
30
+ 3
31
+ >>> segment_tree.right(2)
32
+ 5
33
+ >>> segment_tree.right(12)
34
+ 25
35
+ """
17
36
return idx * 2 + 1
18
37
19
- def build (self , idx , l , r , A ): # noqa: E741
20
- if l == r : # noqa: E741
21
- self .st [idx ] = A [l - 1 ]
38
+ def build (
39
+ self , idx : int , left_element : int , right_element : int , A : List [int ]
40
+ ) -> None :
41
+ if left_element == right_element :
42
+ self .st [idx ] = A [left_element - 1 ]
22
43
else :
23
- mid = (l + r ) // 2
24
- self .build (self .left (idx ), l , mid , A )
25
- self .build (self .right (idx ), mid + 1 , r , A )
44
+ mid = (left_element + right_element ) // 2
45
+ self .build (self .left (idx ), left_element , mid , A )
46
+ self .build (self .right (idx ), mid + 1 , right_element , A )
26
47
self .st [idx ] = max (self .st [self .left (idx )], self .st [self .right (idx )])
27
48
28
- # update with O(lg N) (Normal segment tree without lazy update will take O(Nlg N)
29
- # for each update)
30
- def update ( self , idx , l , r , a , b , val ): # noqa: E741
49
+ def update (
50
+ self , idx : int , left_element : int , right_element : int , a : int , b : int , val : int
51
+ ) -> bool :
31
52
"""
53
+ update with O(lg N) (Normal segment tree without lazy update will take O(Nlg N)
54
+ for each update)
55
+
32
56
update(1, 1, N, a, b, v) for update val v to [a,b]
33
57
"""
34
58
if self .flag [idx ] is True :
35
59
self .st [idx ] = self .lazy [idx ]
36
60
self .flag [idx ] = False
37
- if l != r : # noqa: E741
61
+ if left_element != right_element :
38
62
self .lazy [self .left (idx )] = self .lazy [idx ]
39
63
self .lazy [self .right (idx )] = self .lazy [idx ]
40
64
self .flag [self .left (idx )] = True
41
65
self .flag [self .right (idx )] = True
42
66
43
- if r < a or l > b :
67
+ if right_element < a or left_element > b :
44
68
return True
45
- if l >= a and r <= b : # noqa: E741
69
+ if left_element >= a and right_element <= b :
46
70
self .st [idx ] = val
47
- if l != r : # noqa: E741
71
+ if left_element != right_element :
48
72
self .lazy [self .left (idx )] = val
49
73
self .lazy [self .right (idx )] = val
50
74
self .flag [self .left (idx )] = True
51
75
self .flag [self .right (idx )] = True
52
76
return True
53
- mid = (l + r ) // 2
54
- self .update (self .left (idx ), l , mid , a , b , val )
55
- self .update (self .right (idx ), mid + 1 , r , a , b , val )
77
+ mid = (left_element + right_element ) // 2
78
+ self .update (self .left (idx ), left_element , mid , a , b , val )
79
+ self .update (self .right (idx ), mid + 1 , right_element , a , b , val )
56
80
self .st [idx ] = max (self .st [self .left (idx )], self .st [self .right (idx )])
57
81
return True
58
82
59
83
# query with O(lg N)
60
- def query (self , idx , l , r , a , b ): # noqa: E741
84
+ def query (
85
+ self , idx : int , left_element : int , right_element : int , a : int , b : int
86
+ ) -> int :
61
87
"""
62
88
query(1, 1, N, a, b) for query max of [a,b]
89
+ >>> A = [1, 2, -4, 7, 3, -5, 6, 11, -20, 9, 14, 15, 5, 2, -8]
90
+ >>> segment_tree = SegmentTree(15)
91
+ >>> segment_tree.build(1, 1, 15, A)
92
+ >>> segment_tree.query(1, 1, 15, 4, 6)
93
+ 7
94
+ >>> segment_tree.query(1, 1, 15, 7, 11)
95
+ 14
96
+ >>> segment_tree.query(1, 1, 15, 7, 12)
97
+ 15
63
98
"""
64
99
if self .flag [idx ] is True :
65
100
self .st [idx ] = self .lazy [idx ]
66
101
self .flag [idx ] = False
67
- if l != r : # noqa: E741
102
+ if left_element != right_element :
68
103
self .lazy [self .left (idx )] = self .lazy [idx ]
69
104
self .lazy [self .right (idx )] = self .lazy [idx ]
70
105
self .flag [self .left (idx )] = True
71
106
self .flag [self .right (idx )] = True
72
- if r < a or l > b :
107
+ if right_element < a or left_element > b :
73
108
return - math .inf
74
- if l >= a and r <= b : # noqa: E741
109
+ if left_element >= a and right_element <= b :
75
110
return self .st [idx ]
76
- mid = (l + r ) // 2
77
- q1 = self .query (self .left (idx ), l , mid , a , b )
78
- q2 = self .query (self .right (idx ), mid + 1 , r , a , b )
111
+ mid = (left_element + right_element ) // 2
112
+ q1 = self .query (self .left (idx ), left_element , mid , a , b )
113
+ q2 = self .query (self .right (idx ), mid + 1 , right_element , a , b )
79
114
return max (q1 , q2 )
80
115
81
- def showData (self ):
116
+ def show_data (self ) -> None :
82
117
showList = []
83
118
for i in range (1 , N + 1 ):
84
119
showList += [self .query (1 , 1 , self .N , i , i )]
@@ -96,4 +131,4 @@ def showData(self):
96
131
segt .update (1 , 1 , N , 1 , 3 , 111 )
97
132
print (segt .query (1 , 1 , N , 1 , 15 ))
98
133
segt .update (1 , 1 , N , 7 , 8 , 235 )
99
- segt .showData ()
134
+ segt .show_data ()
0 commit comments