Prefix Sum Array in Python using Accumulate Function

A prefix sum array is an array where each element is the cumulative sum of all the previous elements up to that point in the original array.

For example, given an array [1, 2, 3, 4], its prefix sum array would be [1, 3, 6, 10]. Let us explore various methods to achieve this.

Using accumulate from itertools

This method utilizes the accumulate() function from the itertools module, which is specifically designed to calculate cumulative results like prefix sums.

# Import accumulate from itertools
from itertools import accumulate

# Input array
a = [1, 2, 3, 4]

# Calculate prefix sum array
prefix_sum = list(accumulate(a))
print(prefix_sum) 

[1, 3, 6, 10]

Explanation:

• accumulate() function computes the cumulative sum of elements in the array.
• We convert the result into a list to display the prefix sum array.

Let’s explore some more methods and see how we can calculate prefix sum array in Python using accumulate function.

Using list comprehension

This method uses list comprehension along with a running total to create the prefix sum array in a concise manner.

# Input array
a = [1, 2, 3, 4]

# Initialize a running total
total = 0

# Calculate prefix sum using list comprehension
prefix_sum = [total := total + num for num in a]

print(prefix_sum) 

[1, 3, 6, 10]

Explanation:

• We use the walrus operator (:=) to update the total while iterating over the array.
• The prefix sum is constructed as a list directly.
• This approach combines the functionality of a loop and inline calculation into a single line.

Using numpy

For numerical data, numpy provides an efficient way to compute prefix sums using its cumsum function for calculating prefix sums.

# Import numpy
import numpy as np

# Input array
a = [1, 2, 3, 4]

# Calculate prefix sum using numpy
prefix_sum = np.cumsum(a)

print(prefix_sum.tolist())

[1, 3, 6, 10]

Explanation:

• cumsum function from numpy computes the cumulative sum efficiently.
• The resulting array is converted to a list for easier readability.

Using for loop

A simple approach is to use a for loop to calculate the prefix sum by iterating through the array and maintaining a running total.

# Input array
a = [1, 2, 3, 4]

# Initialize an empty prefix sum array and a running total
p_sum = []
total = 0

# Calculate prefix sum using a loop
for num in a:
    total += num
    p_sum.append(total)
print(p_sum)  

[1, 3, 6, 10]

Explanation:

• We maintain a running total using a variable, which is updated with each element of the array.
• Each element is added to total, and the result is appended to the prefix sum array.

Using reduce from functools

This method uses the reduce() function to calculate the prefix sum array step by step.

# Import reduce from functools
from functools import reduce

# Input array
a = [1, 2, 3, 4]

# Calculate prefix sum using reduce
p_sum = [reduce(lambda x, y: x + y, a[:i+1]) for i in range(len(a))]

print(p_sum) 

[1, 3, 6, 10]

Explanation:

• reduce() function is used to calculate the cumulative sum for each slice of the array.
• Although functional, this method is less efficient than the others due to repeated slicing.