1
- class Node () :
1
+ class Node :
2
2
def __init__ (self , start , end ):
3
3
# Initializes a segment tree node with start and end indices
4
4
self .start = start
@@ -7,17 +7,18 @@ def __init__(self, start, end):
7
7
self .left = None
8
8
self .right = None
9
9
10
- class SegmentTree ():
11
- def __init__ (self , nums , mode = 'max' ):
10
+
11
+ class SegmentTree :
12
+ def __init__ (self , nums , mode = "max" ):
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 )
@@ -38,14 +39,14 @@ def build(self, start, end, nums):
38
39
n = Node (start , end )
39
40
n .value = nums [start ]
40
41
return n
41
-
42
+
42
43
mid = (start + end ) // 2
43
44
root = Node (start , end )
44
45
root .left = self .build (start , mid , nums )
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,24 +58,24 @@ def max_in_range(self, start_index, end_index):
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' )
62
-
61
+ if self .mode == " sum" :
62
+ raise Exception (" Current Segment Tree doesn't support finding max" )
63
+
63
64
if start_index > end_index or start_index < 0 or end_index >= self .size :
64
- raise Exception (' Invalid index' )
65
-
65
+ raise Exception (" Invalid index" )
66
+
66
67
return self .query (self .root , start_index , end_index , 0 , self .size - 1 )
67
68
68
69
def sum_in_range (self , start_index , end_index ):
69
70
"""
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' )
75
-
74
+ if self .mode == " max" :
75
+ raise Exception (" Current Segment Tree doesn't support summing" )
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
@@ -91,7 +92,7 @@ def query(self, node, start_index, end_index, start, end):
91
92
# Complete overlap
92
93
if start_index <= start and end <= end_index :
93
94
return node .value
94
-
95
+
95
96
mid = (start + end ) // 2
96
97
97
98
if end_index <= mid :
@@ -102,10 +103,15 @@ def query(self, node, start_index, end_index, start, end):
102
103
return self .query (node .right , start_index , end_index , mid + 1 , end )
103
104
else :
104
105
# 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 ))
106
+ if self .mode == "max" :
107
+ return max (
108
+ self .query (node .left , start_index , end_index , start , mid ),
109
+ self .query (node .right , start_index , end_index , mid + 1 , end ),
110
+ )
107
111
else :
108
- return (self .query (node .left , start_index , end_index , start , mid ) + self .query (node .right , start_index , end_index , mid + 1 , end ))
112
+ return self .query (
113
+ node .left , start_index , end_index , start , mid
114
+ ) + self .query (node .right , start_index , end_index , mid + 1 , end )
109
115
110
116
def update (self , index , new_value ):
111
117
"""
@@ -114,8 +120,8 @@ def update(self, index, new_value):
114
120
:param new_value: New value to set.
115
121
"""
116
122
if index < 0 or index >= self .size :
117
- raise Exception (' Invalid index' )
118
-
123
+ raise Exception (" Invalid index" )
124
+
119
125
self .modify (self .root , index , new_value , 0 , self .size - 1 )
120
126
121
127
def modify (self , node , index , new_value , start , end ):
@@ -137,9 +143,9 @@ def modify(self, node, index, new_value, start, end):
137
143
self .modify (node .left , index , new_value , start , mid )
138
144
else :
139
145
self .modify (node .right , index , new_value , mid + 1 , end )
140
-
146
+
141
147
# Recompute current node's value after update
142
- if self .mode == ' max' :
148
+ if self .mode == " max" :
143
149
node .value = max (node .left .value , node .right .value )
144
150
else :
145
- node .value = node .left .value + node .right .value
151
+ node .value = node .left .value + node .right .value
0 commit comments