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.
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
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:
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
accumulate is implemented in C and is faster than a Python loop.
NumPy Cumulative Sum
Example: Count Subarrays with Given Sum
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)
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)
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 fromltorisP[r+1] - P[l]. Without the leading zero, it isP[r] - P[l-1]with a special case forl = 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 longorint64. - 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.accumulateornumpy.cumsumfor 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
- Python - Speed up an A Star Pathfinding Algorithm
- Python - Tree traversal question
- Python and OpenCV - Improving my lane detection algorithm
- Python Brute Force algorithm
- Python - Single thread executor already being used, would deadlock
- Python - sklearn How to pass parameters to the customize ModelTransformer class by gridsearchcv
- Python CMA-ES Algorithm to solve user-defined function and constraints
- Python data structure sort list alphabetically

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 courseTrack 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.