Skip to content

Commit cb5762c

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

File tree

1 file changed

+28
-18
lines changed

1 file changed

+28
-18
lines changed

data_structures/binary_tree/segment_tree_node.py

+28-18
Original file line numberDiff line numberDiff line change
@@ -7,17 +7,18 @@ def __init__(self, start: int, end: int) -> None:
77
self.left = None
88
self.right = None
99

10+
1011
class SegmentTree:
11-
def __init__(self, nums: list[int], mode: str='max') -> None:
12+
def __init__(self, nums: list[int], mode: str = "max") -> None:
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)
@@ -45,7 +46,7 @@ def build(self, start: int, end: int, nums: list[int]) -> Node:
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,11 +58,11 @@ def max_in_range(self, start_index: int, end_index: int) -> int:
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')
61+
if self.mode == "sum":
62+
raise Exception("Current Segment Tree doesn't support finding max")
6263

6364
if start_index > end_index or start_index < 0 or end_index >= self.size:
64-
raise Exception('Invalid index')
65+
raise Exception("Invalid index")
6566

6667
return self.query(self.root, start_index, end_index, 0, self.size - 1)
6768

@@ -70,15 +71,17 @@ def sum_in_range(self, start_index: int, end_index: int) -> int:
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')
74+
if self.mode == "max":
75+
raise Exception("Current Segment Tree doesn't support summing")
7576

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

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:
8285
"""
8386
Recursively queries a value (max or sum) in a given range.
8487
: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
102105
return self.query(node.right, start_index, end_index, mid + 1, end)
103106
else:
104107
# 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+
)
107113
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)
109117

110118
def update(self, index: int, new_value: int) -> int:
111119
"""
@@ -114,11 +122,13 @@ def update(self, index: int, new_value: int) -> int:
114122
:param new_value: New value to set.
115123
"""
116124
if index < 0 or index >= self.size:
117-
raise Exception('Invalid index')
125+
raise Exception("Invalid index")
118126

119127
self.modify(self.root, index, new_value, 0, self.size - 1)
120128

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:
122132
"""
123133
Recursively updates the tree to reflect a change at a specific index.
124134
:param node: Current node being processed.
@@ -139,7 +149,7 @@ def modify(self, node: Node, index: int, new_value: int, start: int, end: int) -
139149
self.modify(node.right, index, new_value, mid + 1, end)
140150

141151
# Recompute current node's value after update
142-
if self.mode == 'max':
152+
if self.mode == "max":
143153
node.value = max(node.left.value, node.right.value)
144154
else:
145-
node.value = node.left.value + node.right.value
155+
node.value = node.left.value + node.right.value

0 commit comments

Comments
 (0)