@@ -7,17 +7,18 @@ def __init__(self, start: int, end: int) -> None:
7
7
self .left = None
8
8
self .right = None
9
9
10
+
10
11
class SegmentTree :
11
- def __init__ (self , nums : list [int ], mode : str = ' max' ) -> None :
12
+ def __init__ (self , nums : list [int ], mode : str = " max" ) -> None :
12
13
"""
13
14
Initializes the Segment Tree.
14
15
:param nums: List of integers to build the tree from.
15
16
:param mode: Operation mode of the tree ('max' or 'sum').
16
17
"""
17
18
self .size = len (nums )
18
19
self .mode = mode
19
- if mode not in {' max' , ' sum' }:
20
- self .mode = ' max' # Default to max if invalid mode is given
20
+ if mode not in {" max" , " sum" }:
21
+ self .mode = " max" # Default to max if invalid mode is given
21
22
22
23
# Build the tree from the input list
23
24
self .root = self .build (0 , self .size - 1 , nums )
@@ -45,7 +46,7 @@ def build(self, start: int, end: int, nums: list[int]) -> Node:
45
46
root .right = self .build (mid + 1 , end , nums )
46
47
47
48
# Set the value according to the mode
48
- if self .mode == ' max' :
49
+ if self .mode == " max" :
49
50
root .value = max (root .left .value , root .right .value )
50
51
else :
51
52
root .value = root .left .value + root .right .value
@@ -57,11 +58,11 @@ def max_in_range(self, start_index: int, end_index: int) -> int:
57
58
Queries the maximum value in a given range.
58
59
Only works in 'max' mode.
59
60
"""
60
- if self .mode == ' sum' :
61
- raise Exception (' Current Segment Tree doesn\ ' t support finding max' )
61
+ if self .mode == " sum" :
62
+ raise Exception (" Current Segment Tree doesn't support finding max" )
62
63
63
64
if start_index > end_index or start_index < 0 or end_index >= self .size :
64
- raise Exception (' Invalid index' )
65
+ raise Exception (" Invalid index" )
65
66
66
67
return self .query (self .root , start_index , end_index , 0 , self .size - 1 )
67
68
@@ -70,15 +71,17 @@ def sum_in_range(self, start_index: int, end_index: int) -> int:
70
71
Queries the sum of values in a given range.
71
72
Only works in 'sum' mode.
72
73
"""
73
- if self .mode == ' max' :
74
- raise Exception (' Current Segment Tree doesn\ ' t support summing' )
74
+ if self .mode == " max" :
75
+ raise Exception (" Current Segment Tree doesn't support summing" )
75
76
76
77
if start_index > end_index or start_index < 0 or end_index >= self .size :
77
- raise Exception (' Invalid index' )
78
+ raise Exception (" Invalid index" )
78
79
79
80
return self .query (self .root , start_index , end_index , 0 , self .size - 1 )
80
81
81
- def query (self , node : Node , start_index : int , end_index : int , start : int , end : int ) -> int :
82
+ def query (
83
+ self , node : Node , start_index : int , end_index : int , start : int , end : int
84
+ ) -> int :
82
85
"""
83
86
Recursively queries a value (max or sum) in a given range.
84
87
:param node: Current node in the tree.
@@ -102,10 +105,15 @@ def query(self, node: Node, start_index: int, end_index: int, start: int, end: i
102
105
return self .query (node .right , start_index , end_index , mid + 1 , end )
103
106
else :
104
107
# Range spans both children
105
- if self .mode == 'max' :
106
- return max (self .query (node .left , start_index , end_index , start , mid ), self .query (node .right , start_index , end_index , mid + 1 , end ))
108
+ if self .mode == "max" :
109
+ return max (
110
+ self .query (node .left , start_index , end_index , start , mid ),
111
+ self .query (node .right , start_index , end_index , mid + 1 , end ),
112
+ )
107
113
else :
108
- return self .query (node .left , start_index , end_index , start , mid ) + self .query (node .right , start_index , end_index , mid + 1 , end )
114
+ return self .query (
115
+ node .left , start_index , end_index , start , mid
116
+ ) + self .query (node .right , start_index , end_index , mid + 1 , end )
109
117
110
118
def update (self , index : int , new_value : int ) -> int :
111
119
"""
@@ -114,11 +122,13 @@ def update(self, index: int, new_value: int) -> int:
114
122
:param new_value: New value to set.
115
123
"""
116
124
if index < 0 or index >= self .size :
117
- raise Exception (' Invalid index' )
125
+ raise Exception (" Invalid index" )
118
126
119
127
self .modify (self .root , index , new_value , 0 , self .size - 1 )
120
128
121
- def modify (self , node : Node , index : int , new_value : int , start : int , end : int ) -> int :
129
+ def modify (
130
+ self , node : Node , index : int , new_value : int , start : int , end : int
131
+ ) -> int :
122
132
"""
123
133
Recursively updates the tree to reflect a change at a specific index.
124
134
:param node: Current node being processed.
@@ -139,7 +149,7 @@ def modify(self, node: Node, index: int, new_value: int, start: int, end: int) -
139
149
self .modify (node .right , index , new_value , mid + 1 , end )
140
150
141
151
# Recompute current node's value after update
142
- if self .mode == ' max' :
152
+ if self .mode == " max" :
143
153
node .value = max (node .left .value , node .right .value )
144
154
else :
145
- node .value = node .left .value + node .right .value
155
+ node .value = node .left .value + node .right .value
0 commit comments