Prefix Sums

Resources

Introduction

Let's say we have a one-indexed integer array $\text{arr}$ of size $N$ and we want to compute the value of

for $Q$ different pairs $(a,b)$ satisfying $1\le a\le b\le N$. We'll use the following example with $N=6$:

Naively, for every query, we can iterate through all entries from index $a$ to index $b$ to add them up. Since we have $Q$ queries and each query requires a maximum of $\mathcal{O}(N)$ operations to calculate the sum, our total time complexity is $\mathcal{O}(NQ)$. For most problems of this nature, the constraints will be $N,Q\le 10^5$, so $NQ$ is on the order of $10^{10}$. This is not acceptable; it will almost certainly exceed the time limit.

Instead, we can use prefix sums to process these array sum queries. We designate a prefix sum array $\text{prefix}$. First, because we're 1-indexing the array, set $\text{prefix}[0]=0$, then for indices $k$ such that $1\le k\le n$, define the prefix sum array as follows:

Basically, what this means is that the element at index $k$ of the prefix sum array stores the sum of all the elements in the original array from index $1$ up to $k$. This can be calculated easily in $\mathcal{O}(N)$ by the following formula for each $1\le k\le n$:

For the example case, our prefix sum array looks like this:

Now, when we want to query for the sum of the elements of $\text{arr}$ between (1-indexed) indices $a$ and $b$ inclusive, we can use the following formula:

Using our definition of the elements in the prefix sum array, we have

Since we are only querying two elements in the prefix sum array, we can calculate subarray sums in $\mathcal{O}(1)$ per query, which is much better than the $\mathcal{O}(N)$ per query that we had before. Now, after an $\mathcal{O}(N)$ preprocessing to calculate the prefix sum array, each of the $Q$ queries takes $\mathcal{O}(1)$ time. Thus, our total time complexity is $\mathcal{O}(N+Q)$, which should now pass the time limit.

Let's do an example query and find the subarray sum between indices $a=2$ and $b=5$, inclusive, in the 1-indexed $\text{arr}$. From looking at the original array, we see that this is

Using prefix sums:

These are also known as partial sums.

Code

1#include <bits/stdc++.h>
2using namespace std;
3 
4#define sz(x) (int)size(x)
5
6using ll = long long;
7using vl = vector<ll>;
8
9vl psum(const vl& a) {
10    vl psum(sz(a)+1); 

Problems

Max Subarray Sum

Now we'll look at some extensions.

This problem has a solution known as Kadane's Algorithm. Please don't use that solution; try to solve it with prefix sums.

Prefix Minimum, XOR, etc.

Similar to prefix sums, you can also take prefix minimum or maximum; but you cannot answer min queries over an arbitrary range with prefix minimum. (This is because minimum doesn't have an inverse operation, the way subtraction is to addition.) On the other hand, XOR is its own inverse operation, meaning that the XOR of any number with itself is zero.

2D Prefix Sums

Now, what if we wanted to process $Q$ queries for the sum over a subrectangle of a $N$ rows by $M$ columns matrix in two dimensions? Let's assume both rows and columns are 1-indexed, and we use the following matrix as an example:

Naively, each sum query would then take $\mathcal{O}(NM)$ time, for a total of $\mathcal{O}(QNM)$. This is too slow.

Let's take the following example region, which we want to sum:

Manually summing all the cells, we have a submatrix sum of $7+11+9+6+1+3=37$.

The first logical optimization would be to do one-dimensional prefix sums of each row. Then, we'd have the following row-prefix sum matrix. The desired subarray sum of each row in our desired region is simply the green cell minus the red cell in that respective row. We do this for each row to get $(28-1)+(14-4)=37$.

Now, if we wanted to find a submatrix sum, we could break up the submatrix into a subarray for each row, and then add their sums, which would be calculated using the prefix sums method described earlier. Since the matrix has $N$ rows, the time complexity of this is $\mathcal{O}(QN)$. This might be fast enough for $Q=10^5$ and $N=10^3$, but we can do better.

In fact, we can do two-dimensional prefix sums. In our two dimensional prefix sum array, we have

This can be calculated as follows for row index $1\le i\le n$ and column index $1\le j\le m$:

The submatrix sum between rows $a$ and $A$ and columns $b$ and $B$, can thus be expressed as follows:

Summing the blue region from above using the 2D prefix sums method, we add the value of the green square, subtract the values of the red squares, and then add the value of the gray square. In this example, we have

as expected.

Since no matter the size of the submatrix we are summing, we only need to access four values of the 2D prefix sum array, this runs in $\mathcal{O}(1)$ per query after an $\mathcal{O}(NM)$ preprocessing.

More Complex Applications

Instead of storing just the values themselves, you can also take a prefix sum over $i\cdot a_i$, or $10^i\cdot a_i$, for instance. Some math is usually helpful for these problems; don't be afraid to get dirty with algebra!

For instance, we can quickly answer the following type of query:

Find $1\cdot a_l+2\cdot a_{l+1}+3\cdot a_{l+2}+\cdots +(r-l+1)\cdot a_r$.

First, define the following:

Then, we have:

And so,

Which is what we were looking for!