add new file data_structures/binary_tree/segment_tree_node.py add ret type

Author: xxx595588

data_structures/binary_tree/segment_tree_node.py

@@ -16,8 +16,8 @@ def __init__(self, nums: list[int], mode: str='max') -> None:
         """
         self.size = len(nums)
         self.mode = mode
-        if mode not in {'max', 'sum'}:
-            self.mode = 'max'  # Default to max if invalid mode is given
+        if mode not in {"max", "sum"}:
+            self.mode = "max"  # Default to max if invalid mode is given
 
         # Build the tree from the input list
         self.root = self.build(0, self.size - 1, nums)
@@ -38,14 +38,14 @@ def build(self, start: int, end: int, nums: list[int]) -> Node:
             n = Node(start, end)
             n.value = nums[start]
             return n
-        
+
         mid = (start + end) // 2
         root = Node(start, end)
         root.left = self.build(start, mid, nums)
         root.right = self.build(mid + 1, end, nums)
 
         # Set the value according to the mode
-        if self.mode == 'max':
+        if self.mode == "max":
             root.value = max(root.left.value, root.right.value)
         else:
             root.value = root.left.value + root.right.value
@@ -57,24 +57,24 @@ def max_in_range(self, start_index: int, end_index: int) -> int:
         Queries the maximum value in a given range.
         Only works in 'max' mode.
         """
-        if self.mode == 'sum':
-            raise Exception('Current Segment Tree doesn\'t support finding max') 
-        
+        if self.mode == "sum":
+            raise Exception("Current Segment Tree doesn't support finding max")
+
         if start_index > end_index or start_index < 0 or end_index >= self.size:
-            raise Exception('Invalid index')
-        
+            raise Exception("Invalid index")
+
         return self.query(self.root, start_index, end_index, 0, self.size - 1)
 
     def sum_in_range(self, start_index: int, end_index: int) -> int:
         """
         Queries the sum of values in a given range.
         Only works in 'sum' mode.
         """
-        if self.mode == 'max':
-            raise Exception('Current Segment Tree doesn\'t support summing') 
-        
+        if self.mode == "max":
+            raise Exception("Current Segment Tree doesn't support summing")
+
         if start_index > end_index or start_index < 0 or end_index >= self.size:
-            raise Exception('Invalid index')
+            raise Exception("Invalid index")
 
         return self.query(self.root, start_index, end_index, 0, self.size - 1)
 
@@ -91,7 +91,7 @@ def query(self, node: Node, start_index: int, end_index: int, start: int, end: i
         # Complete overlap
         if start_index <= start and end <= end_index:
             return node.value
-        
+
         mid = (start + end) // 2
 
         if end_index <= mid:
@@ -102,10 +102,15 @@ def query(self, node: Node, start_index: int, end_index: int, start: int, end: i
             return self.query(node.right, start_index, end_index, mid + 1, end)
         else:
             # Range spans both children
-            if self.mode == 'max':
-                return max(self.query(node.left, start_index, end_index, start, mid), self.query(node.right, start_index, end_index, mid + 1, end))
+            if self.mode == "max":
+                return max(
+                    self.query(node.left, start_index, end_index, start, mid),
+                    self.query(node.right, start_index, end_index, mid + 1, end),
+                )
             else:
-                return (self.query(node.left, start_index, end_index, start, mid) + self.query(node.right, start_index, end_index, mid + 1, end))
+                return self.query(
+                    node.left, start_index, end_index, start, mid
+                ) + self.query(node.right, start_index, end_index, mid + 1, end)
 
     def update(self, index: int, new_value: int) -> int:
         """
@@ -114,8 +119,8 @@ def update(self, index: int, new_value: int) -> int:
         :param new_value: New value to set.
         """
         if index < 0 or index >= self.size:
-            raise Exception('Invalid index')
-        
+            raise Exception("Invalid index")
+
         self.modify(self.root, index, new_value, 0, self.size - 1)
 
     def modify(self, node: Node, index: int, new_value: int, start: int, end: int) -> int:
@@ -137,9 +142,9 @@ def modify(self, node: Node, index: int, new_value: int, start: int, end: int) -
             self.modify(node.left, index, new_value, start, mid)
         else:
             self.modify(node.right, index, new_value, mid + 1, end)
-        
+
         # Recompute current node's value after update
-        if self.mode == 'max':
+        if self.mode == "max":
             node.value = max(node.left.value, node.right.value)
         else:
             node.value = node.left.value + node.right.value