Python
Prefix Sum
Algorithm
Coding
Programming

python - prefix sum algorithm

Data Structures & Algorithms practice on Codemia

Step through 300 algorithm problems with animated visualisers that show the data structure changing as the code runs.

Practice algorithms

Introduction

A prefix sum (cumulative sum) array stores the running total of elements up to each index. With a prefix sum array pre-computed in O(n), you can answer any range sum query sum(a[l:r+1]) in O(1) instead of O(n). This technique is foundational in competitive programming, database query optimization, and image processing (integral images).

Building a Prefix Sum Array

python
1def prefix_sum(arr):
2    n = len(arr)
3    P = [0] * (n + 1)  # P[0] = 0, so range queries work cleanly
4    for i in range(n):
5        P[i + 1] = P[i] + arr[i]
6    return P
7
8arr = [3, 1, 4, 1, 5, 9, 2, 6]
9P = prefix_sum(arr)
10print(P)  # [0, 3, 4, 8, 9, 14, 23, 25, 31]

P[i] stores the sum of arr[0:i]. Using a 1-indexed prefix array (with P[0] = 0) simplifies range queries.

Range Sum Queries

The sum of elements from index l to r (inclusive) is:

python
1def range_sum(P, l, r):
2    return P[r + 1] - P[l]
3
4arr = [3, 1, 4, 1, 5, 9, 2, 6]
5P = prefix_sum(arr)
6
7# Sum of arr[2:5] = 4 + 1 + 5 + 9 = 19
8print(range_sum(P, 2, 5))  # 19
9
10# Sum of arr[0:3] = 3 + 1 + 4 + 1 = 9
11print(range_sum(P, 0, 3))  # 9
12
13# Sum of entire array
14print(range_sum(P, 0, len(arr) - 1))  # 31

Without prefix sums, each range query takes O(n). With prefix sums, precomputation is O(n) and each query is O(1).

Using itertools.accumulate

python
1from itertools import accumulate
2
3arr = [3, 1, 4, 1, 5, 9, 2, 6]
4P = [0] + list(accumulate(arr))
5print(P)  # [0, 3, 4, 8, 9, 14, 23, 25, 31]
6
7# Range sum query: sum of arr[2:5]
8print(P[6] - P[2])  # 19

accumulate is implemented in C and is faster than a Python loop.

NumPy Cumulative Sum

python
1import numpy as np
2
3arr = np.array([3, 1, 4, 1, 5, 9, 2, 6])
4P = np.concatenate(([0], np.cumsum(arr)))
5print(P)  # [ 0  3  4  8  9 14 23 25 31]
6
7# Vectorized range queries
8starts = np.array([0, 2, 4])
9ends = np.array([3, 5, 7])
10sums = P[ends + 1] - P[starts]
11print(sums)  # [ 9 19 22]

Example: Count Subarrays with Given Sum

python
1from collections import defaultdict
2
3def count_subarrays_with_sum(arr, target):
4    """Count subarrays that sum to target. O(n) time."""
5    count = 0
6    current_sum = 0
7    prefix_counts = defaultdict(int)
8    prefix_counts[0] = 1  # Empty prefix
9
10    for num in arr:
11        current_sum += num
12        # If (current_sum - target) was a previous prefix sum,
13        # then the subarray between them sums to target
14        count += prefix_counts[current_sum - target]
15        prefix_counts[current_sum] += 1
16
17    return count
18
19arr = [1, 2, 3, -2, 5]
20print(count_subarrays_with_sum(arr, 3))  # 3: [3], [1,2], [-2,5]

This uses the prefix sum concept without building the full array — it checks if a complementary prefix sum exists using a hash map.

2D Prefix Sum (Integral Image)

python
1def prefix_sum_2d(matrix):
2    rows, cols = len(matrix), len(matrix[0])
3    P = [[0] * (cols + 1) for _ in range(rows + 1)]
4
5    for r in range(rows):
6        for c in range(cols):
7            P[r+1][c+1] = (matrix[r][c]
8                          + P[r][c+1]
9                          + P[r+1][c]
10                          - P[r][c])
11    return P
12
13def region_sum(P, r1, c1, r2, c2):
14    """Sum of submatrix from (r1,c1) to (r2,c2) inclusive."""
15    return (P[r2+1][c2+1]
16           - P[r1][c2+1]
17           - P[r2+1][c1]
18           + P[r1][c1])
19
20matrix = [
21    [1, 2, 3],
22    [4, 5, 6],
23    [7, 8, 9]
24]
25P = prefix_sum_2d(matrix)
26
27# Sum of submatrix (1,1) to (2,2) = 5+6+8+9 = 28
28print(region_sum(P, 1, 1, 2, 2))  # 28
29
30# Sum of entire matrix = 45
31print(region_sum(P, 0, 0, 2, 2))  # 45

2D prefix sums enable O(1) rectangular region queries after O(rows * cols) precomputation. This is called an "integral image" in computer vision.

Difference Array (Inverse of Prefix Sum)

python
1def apply_range_updates(n, updates):
2    """Apply multiple range increment operations efficiently."""
3    diff = [0] * (n + 1)
4
5    for l, r, val in updates:
6        diff[l] += val
7        diff[r + 1] -= val
8
9    # Prefix sum of difference array gives final values
10    result = [0] * n
11    result[0] = diff[0]
12    for i in range(1, n):
13        result[i] = result[i-1] + diff[i]
14
15    return result
16
17# Add 3 to indices 1-4, add 5 to indices 2-6
18updates = [(1, 4, 3), (2, 6, 5)]
19print(apply_range_updates(8, updates))
20# [0, 3, 8, 8, 8, 5, 5, 0]

The difference array is the inverse of prefix sum. It applies O(k) range updates in O(n + k) total time instead of O(n * k).

Common Pitfalls

  • Off-by-one errors: Using 0-indexed vs 1-indexed prefix arrays changes the range query formula. With P[0] = 0, the sum from l to r is P[r+1] - P[l]. Without the leading zero, it is P[r] - P[l-1] with a special case for l = 0.
  • Integer overflow: Large arrays with large values can overflow 32-bit integers. Python handles arbitrary precision natively, but in C/C++/Java, use long long or int64.
  • Modifying the original array: Prefix sums assume the array is static. If elements change, the entire prefix array must be recomputed. For dynamic arrays, use a Fenwick tree (Binary Indexed Tree) instead.
  • 2D formula sign errors: The inclusion-exclusion formula P[r2][c2] - P[r1-1][c2] - P[r2][c1-1] + P[r1-1][c1-1] has four terms. Getting the signs wrong gives incorrect results.
  • Using prefix sums for min/max: Prefix sums only work for sum queries. For range minimum/maximum queries, use a sparse table or segment tree instead.

Summary

  • Build a prefix sum array in O(n) to answer range sum queries in O(1)
  • Use itertools.accumulate or numpy.cumsum for clean implementations
  • The range sum formula is P[r+1] - P[l] (with a leading zero in the prefix array)
  • 2D prefix sums enable O(1) rectangular region queries (integral images)
  • Difference arrays are the inverse — they efficiently apply range updates
  • For dynamic arrays with updates, use a Fenwick tree instead of rebuilding prefix sums

Related reading
Course
Intermediate
27 lessons
15 hours
DSA Fundamentals

Master algorithmic patterns and data structures through hands-on LeetCode-style problems - from arrays and hashing to dynamic programming and advanced graphs.

View the course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

Data Structures & Algorithms practice on Codemia

Step through 300 algorithm problems with animated visualisers that show the data structure changing as the code runs.

Practice algorithms

All Rights Reserved.