Skip to content

Commit f4fe48b

Browse files
authored
Create prefix_sum.py
1 parent 183fa06 commit f4fe48b

File tree

1 file changed

+90
-0
lines changed

1 file changed

+90
-0
lines changed

dynamic_programming/prefix_sum.py

+90
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
1+
"""
2+
Author : Sanjay Muthu <https://github.com/XenoBytesX>
3+
4+
This is an implementation of the Dynamic Programming solution to the Range Sum Query problem.
5+
6+
The problem statement is:
7+
Given an array and q queries,
8+
each query stating you to find the sum of elements l to r (inclusive) (l and r are given in the query)
9+
10+
Example:
11+
arr = [1, 4, 6, 2, 61, 12]
12+
queries = 3
13+
l_1 = 2, r_1 = 5
14+
l_2 = 1, r_2 = 5
15+
l_3 = 3, r_3 = 4
16+
17+
as input will return
18+
19+
[81, 85, 63]
20+
21+
as output
22+
23+
[0-indexing]
24+
25+
NOTE:- 0-indexing means the indexing of the array starts from 0
26+
Example:- a = [1, 2, 3, 4, 5, 6]
27+
Here, the 0th index of a is 1,
28+
the 1st index of a is 2,
29+
and so forth
30+
31+
Time Complexity:- O(N + Q)
32+
-- O(N) pre-calculation time to calculate the prefix sum array
33+
-- and O(1) time per each query = O(1 * Q) = O(Q) time
34+
35+
Space Complexity:- O(N)
36+
-- O(N) to store the prefix sum
37+
38+
Algorithm:-
39+
So first we calculate the prefix sum (dp) of the array.
40+
The prefix sum of the index i is the sum of all elements indexed from 0 to i (inclusive)
41+
The prefix sum of the index i is the prefix sum of index i-1 + the current element.
42+
So, The state of the dp is dp[i] = dp[i-1] + a[i];
43+
44+
After we calculate the prefix sum,
45+
For each query [l, r]
46+
The answer is dp[r]-dp[l-1] ( we need to be careful because l might be 0 )
47+
For example take this array:-
48+
[4, 2, 1, 6, 3]
49+
The prefix sum calculated for this array would be:-
50+
[4, 4+2, 4+2+1, 4+2+1+6, 4+2+1+6+3]
51+
==> [4, 6, 7, 13, 16]
52+
If the query was l=3, r=4,
53+
The answer would be 6+3 = 9 but this would require O(r-l) time ≈ O(N) time
54+
If we use prefix sums we can find it in O(1) time by using the formula prefix[r]-prefix[l-1]
55+
This formula works because prefix[r] is the sum of elements from [0, r] and prefix[l-1] is the sum of elements from [0, l-1],
56+
so if we do prefix[r]-prefix[l-1] it will be [0, r] - [0, l-1] = [0, l-1] + [l, r] - [0, l-1] = [l, r]
57+
"""
58+
59+
from __future__ import annotations
60+
61+
62+
def prefix_sum(array: list[int], queries: list[tuple[int, int]]) -> list[int]:
63+
"""
64+
Some examples
65+
>>> prefix_sum([1, 4, 6, 2, 61, 12], [(2, 5), (1, 5), (3, 4)])
66+
[81, 85, 63]
67+
>>> prefix_sum([4, 2, 1, 6, 3], [(3, 4), (1, 3), (0, 2)])
68+
[9, 9, 7]
69+
"""
70+
# The prefix sum array
71+
dp = [0] * len(array)
72+
dp[0] = array[0]
73+
for i in range(1, len(array)):
74+
dp[i] = dp[i-1] + array[i]
75+
76+
# Read Algorithm section (Line 38)
77+
result = []
78+
for query in queries:
79+
res = dp[query[1]]
80+
if query[0] != 0:
81+
res -= dp[query[0]-1]
82+
result.append(res)
83+
84+
return result
85+
86+
87+
if __name__ == "__main__":
88+
import doctest
89+
90+
doctest.testmod()

0 commit comments

Comments
 (0)