Update prefix_sum.py

data_structures/arrays/prefix_sum.py

@@ -1,12 +1,6 @@
-"""
-Author  : Alexander Pantyukhin
-Date    : November 3, 2022
 
 Implement the class of prefix sum with useful functions based on it.
 
-"""
-
-
 class PrefixSum:
     def __init__(self, array: list[int]) -> None:
         len_array = len(array)
@@ -22,7 +16,7 @@
         """
         The function returns the sum of array from the start to the end indexes.
         Runtime : O(1)
-        Space: O(1)
+        Space: O(n) (due to the prefix sum array)
 
         >>> PrefixSum([1,2,3]).get_sum(0, 2)
         6
@@ -35,14 +29,17 @@
         ...
         IndexError: list index out of range
         """
+        if start > end or end >= len(self.prefix_sum):
+            raise IndexError("Invalid range specified.")
+
         if start == 0:
             return self.prefix_sum[end]
 
         return self.prefix_sum[end] - self.prefix_sum[start - 1]
 
     def contains_sum(self, target_sum: int) -> bool:
         """
-        The function returns True if array contains the target_sum,
+        The function returns True if any subarray contains the target_sum,
         False otherwise.
 
         Runtime : O(n)
@@ -64,7 +61,7 @@
 
         sums = {0}
         for sum_item in self.prefix_sum:
-            if sum_item - target_sum in sums:
+            if sum_item == target_sum or sum_item - target_sum in sums:
                 return True
 
             sums.add(sum_item)
@@ -76,3 +73,10 @@
     import doctest
 
     doctest.testmod()
+
+
+
+
+
+
+