33
44
55class SegmentTree :
6- def __init__ (self , N : int ) -> None :
7- self .N = N
8- # approximate the overall size of segment tree with array N
9- self .st : List [ int ] = [0 for i in range (0 , 4 * N )]
6+ def __init__ (self , size : int ) -> None :
7+ self .size = size
8+ # approximate the overall size of segment tree with given value
9+ self .segment_tree = [0 for i in range (0 , 4 * size )]
1010 # 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
11+ self .lazy = [0 for i in range (0 , 4 * size )]
12+ self .flag = [0 for i in range (0 , 4 * size )] # flag for lazy update
1313
1414 def left (self , idx : int ) -> int :
1515 """
@@ -39,24 +39,26 @@ def build(
3939 self , idx : int , left_element : int , right_element : int , A : List [int ]
4040 ) -> None :
4141 if left_element == right_element :
42- self .st [idx ] = A [left_element - 1 ]
42+ self .segment_tree [idx ] = A [left_element - 1 ]
4343 else :
4444 mid = (left_element + right_element ) // 2
4545 self .build (self .left (idx ), left_element , mid , A )
4646 self .build (self .right (idx ), mid + 1 , right_element , A )
47- self .st [idx ] = max (self .st [self .left (idx )], self .st [self .right (idx )])
47+ self .segment_tree [idx ] = max (
48+ self .segment_tree [self .left (idx )], self .segment_tree [self .right (idx )]
49+ )
4850
4951 def update (
5052 self , idx : int , left_element : int , right_element : int , a : int , b : int , val : int
5153 ) -> bool :
5254 """
53- update with O(lg N ) (Normal segment tree without lazy update will take O(Nlg N )
55+ update with O(lg n ) (Normal segment tree without lazy update will take O(nlg n )
5456 for each update)
5557
56- update(1, 1, N , a, b, v) for update val v to [a,b]
58+ update(1, 1, size , a, b, v) for update val v to [a,b]
5759 """
5860 if self .flag [idx ] is True :
59- self .st [idx ] = self .lazy [idx ]
61+ self .segment_tree [idx ] = self .lazy [idx ]
6062 self .flag [idx ] = False
6163 if left_element != right_element :
6264 self .lazy [self .left (idx )] = self .lazy [idx ]
@@ -67,7 +69,7 @@ def update(
6769 if right_element < a or left_element > b :
6870 return True
6971 if left_element >= a and right_element <= b :
70- self .st [idx ] = val
72+ self .segment_tree [idx ] = val
7173 if left_element != right_element :
7274 self .lazy [self .left (idx )] = val
7375 self .lazy [self .right (idx )] = val
@@ -77,15 +79,17 @@ def update(
7779 mid = (left_element + right_element ) // 2
7880 self .update (self .left (idx ), left_element , mid , a , b , val )
7981 self .update (self .right (idx ), mid + 1 , right_element , a , b , val )
80- self .st [idx ] = max (self .st [self .left (idx )], self .st [self .right (idx )])
82+ self .segment_tree [idx ] = max (
83+ self .segment_tree [self .left (idx )], self .segment_tree [self .right (idx )]
84+ )
8185 return True
8286
83- # query with O(lg N )
87+ # query with O(lg n )
8488 def query (
8589 self , idx : int , left_element : int , right_element : int , a : int , b : int
8690 ) -> int :
8791 """
88- query(1, 1, N , a, b) for query max of [a,b]
92+ query(1, 1, size , a, b) for query max of [a,b]
8993 >>> A = [1, 2, -4, 7, 3, -5, 6, 11, -20, 9, 14, 15, 5, 2, -8]
9094 >>> segment_tree = SegmentTree(15)
9195 >>> segment_tree.build(1, 1, 15, A)
@@ -97,7 +101,7 @@ def query(
97101 15
98102 """
99103 if self .flag [idx ] is True :
100- self .st [idx ] = self .lazy [idx ]
104+ self .segment_tree [idx ] = self .lazy [idx ]
101105 self .flag [idx ] = False
102106 if left_element != right_element :
103107 self .lazy [self .left (idx )] = self .lazy [idx ]
@@ -107,28 +111,25 @@ def query(
107111 if right_element < a or left_element > b :
108112 return - math .inf
109113 if left_element >= a and right_element <= b :
110- return self .st [idx ]
114+ return self .segment_tree [idx ]
111115 mid = (left_element + right_element ) // 2
112116 q1 = self .query (self .left (idx ), left_element , mid , a , b )
113117 q2 = self .query (self .right (idx ), mid + 1 , right_element , a , b )
114118 return max (q1 , q2 )
115119
116- def show_data (self ) -> None :
117- showList = []
118- for i in range (1 , N + 1 ):
119- showList += [self .query (1 , 1 , self .N , i , i )]
120- print (showList )
120+ def __str__ (self ) -> None :
121+ return [self .query (1 , 1 , self .size , i , i ) for i in range (1 , self .size + 1 )]
121122
122123
123124if __name__ == "__main__" :
124125 A = [1 , 2 , - 4 , 7 , 3 , - 5 , 6 , 11 , - 20 , 9 , 14 , 15 , 5 , 2 , - 8 ]
125- N = 15
126- segt = SegmentTree (N )
127- segt .build (1 , 1 , N , A )
128- print (segt .query (1 , 1 , N , 4 , 6 ))
129- print (segt .query (1 , 1 , N , 7 , 11 ))
130- print (segt .query (1 , 1 , N , 7 , 12 ))
131- segt .update (1 , 1 , N , 1 , 3 , 111 )
132- print (segt .query (1 , 1 , N , 1 , 15 ))
133- segt .update (1 , 1 , N , 7 , 8 , 235 )
134- segt . show_data ( )
126+ size = 15
127+ segt = SegmentTree (size )
128+ segt .build (1 , 1 , size , A )
129+ print (segt .query (1 , 1 , size , 4 , 6 ))
130+ print (segt .query (1 , 1 , size , 7 , 11 ))
131+ print (segt .query (1 , 1 , size , 7 , 12 ))
132+ segt .update (1 , 1 , size , 1 , 3 , 111 )
133+ print (segt .query (1 , 1 , size , 1 , 15 ))
134+ segt .update (1 , 1 , size , 7 , 8 , 235 )
135+ print ( segt )
0 commit comments