Python/data_structures/arrays/prefix_sum.py
Pranjay kumar 580273eeca
Improve prefix_sum.py (#12560)
* Update prefix_sum.py

Index Validation for get_sum

Raises ValueError if start or end is out of range or start > end.
Handles cases where the array is empty.
 Empty Array Support

If an empty array is passed, get_sum raises an appropriate error instead of failing unexpectedly.
 Optimized contains_sum Initialization

Initializes sums with {0} for efficient subarray sum checking.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update prefix_sum.py

* Update prefix_sum.py

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update prefix_sum.py

* Update prefix_sum.py

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Maxim Smolskiy <mithridatus@mail.ru>
2025-03-20 02:33:46 +03:00

97 lines
2.5 KiB
Python

"""
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)
self.prefix_sum = [0] * len_array
if len_array > 0:
self.prefix_sum[0] = array[0]
for i in range(1, len_array):
self.prefix_sum[i] = self.prefix_sum[i - 1] + array[i]
def get_sum(self, start: int, end: int) -> int:
"""
The function returns the sum of array from the start to the end indexes.
Runtime : O(1)
Space: O(1)
>>> PrefixSum([1,2,3]).get_sum(0, 2)
6
>>> PrefixSum([1,2,3]).get_sum(1, 2)
5
>>> PrefixSum([1,2,3]).get_sum(2, 2)
3
>>> PrefixSum([]).get_sum(0, 0)
Traceback (most recent call last):
...
ValueError: The array is empty.
>>> PrefixSum([1,2,3]).get_sum(-1, 2)
Traceback (most recent call last):
...
ValueError: Invalid range specified.
>>> PrefixSum([1,2,3]).get_sum(2, 3)
Traceback (most recent call last):
...
ValueError: Invalid range specified.
>>> PrefixSum([1,2,3]).get_sum(2, 1)
Traceback (most recent call last):
...
ValueError: Invalid range specified.
"""
if not self.prefix_sum:
raise ValueError("The array is empty.")
if start < 0 or end >= len(self.prefix_sum) or start > end:
raise ValueError("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,
False otherwise.
Runtime : O(n)
Space: O(n)
>>> PrefixSum([1,2,3]).contains_sum(6)
True
>>> PrefixSum([1,2,3]).contains_sum(5)
True
>>> PrefixSum([1,2,3]).contains_sum(3)
True
>>> PrefixSum([1,2,3]).contains_sum(4)
False
>>> PrefixSum([1,2,3]).contains_sum(7)
False
>>> PrefixSum([1,-2,3]).contains_sum(2)
True
"""
sums = {0}
for sum_item in self.prefix_sum:
if sum_item - target_sum in sums:
return True
sums.add(sum_item)
return False
if __name__ == "__main__":
import doctest
doctest.testmod()