Skip to content

Commit cd6dca2

Browse files
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
1 parent 571214a commit cd6dca2

File tree

1 file changed

+31
-25
lines changed

1 file changed

+31
-25
lines changed

data_structures/binary_tree/segment_tree_node.py

+31-25
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
class Node():
1+
class Node:
22
def __init__(self, start, end):
33
# Initializes a segment tree node with start and end indices
44
self.start = start
@@ -7,17 +7,18 @@ def __init__(self, start, end):
77
self.left = None
88
self.right = None
99

10-
class SegmentTree():
11-
def __init__(self, nums, mode='max'):
10+
11+
class SegmentTree:
12+
def __init__(self, nums, mode="max"):
1213
"""
1314
Initializes the Segment Tree.
1415
:param nums: List of integers to build the tree from.
1516
:param mode: Operation mode of the tree ('max' or 'sum').
1617
"""
1718
self.size = len(nums)
1819
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
2122

2223
# Build the tree from the input list
2324
self.root = self.build(0, self.size - 1, nums)
@@ -38,14 +39,14 @@ def build(self, start, end, nums):
3839
n = Node(start, end)
3940
n.value = nums[start]
4041
return n
41-
42+
4243
mid = (start + end) // 2
4344
root = Node(start, end)
4445
root.left = self.build(start, mid, nums)
4546
root.right = self.build(mid + 1, end, nums)
4647

4748
# Set the value according to the mode
48-
if self.mode == 'max':
49+
if self.mode == "max":
4950
root.value = max(root.left.value, root.right.value)
5051
else:
5152
root.value = root.left.value + root.right.value
@@ -57,24 +58,24 @@ def max_in_range(self, start_index, end_index):
5758
Queries the maximum value in a given range.
5859
Only works in 'max' mode.
5960
"""
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+
6364
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+
6667
return self.query(self.root, start_index, end_index, 0, self.size - 1)
6768

6869
def sum_in_range(self, start_index, end_index):
6970
"""
7071
Queries the sum of values in a given range.
7172
Only works in 'sum' mode.
7273
"""
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+
7677
if start_index > end_index or start_index < 0 or end_index >= self.size:
77-
raise Exception('Invalid index')
78+
raise Exception("Invalid index")
7879

7980
return self.query(self.root, start_index, end_index, 0, self.size - 1)
8081

@@ -91,7 +92,7 @@ def query(self, node, start_index, end_index, start, end):
9192
# Complete overlap
9293
if start_index <= start and end <= end_index:
9394
return node.value
94-
95+
9596
mid = (start + end) // 2
9697

9798
if end_index <= mid:
@@ -102,10 +103,15 @@ def query(self, node, start_index, end_index, start, end):
102103
return self.query(node.right, start_index, end_index, mid + 1, end)
103104
else:
104105
# 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+
)
107111
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)
109115

110116
def update(self, index, new_value):
111117
"""
@@ -114,8 +120,8 @@ def update(self, index, new_value):
114120
:param new_value: New value to set.
115121
"""
116122
if index < 0 or index >= self.size:
117-
raise Exception('Invalid index')
118-
123+
raise Exception("Invalid index")
124+
119125
self.modify(self.root, index, new_value, 0, self.size - 1)
120126

121127
def modify(self, node, index, new_value, start, end):
@@ -137,9 +143,9 @@ def modify(self, node, index, new_value, start, end):
137143
self.modify(node.left, index, new_value, start, mid)
138144
else:
139145
self.modify(node.right, index, new_value, mid + 1, end)
140-
146+
141147
# Recompute current node's value after update
142-
if self.mode == 'max':
148+
if self.mode == "max":
143149
node.value = max(node.left.value, node.right.value)
144150
else:
145-
node.value = node.left.value + node.right.value
151+
node.value = node.left.value + node.right.value

0 commit comments

Comments
 (0)