Prefix Sums

One common component of algorithm design is precomputation. That is, taking part of an algorithm that would normally be slow and making it faster by doing some upfront work. That's a little abstract, so let's dive straight in with an example problem.

Mansion (AIO 2007)

Have a read of the problem Mansion from AIO 2007 and spend a few minutes trying the problem for yourself. In summary, you're given an array (for python users: an array is another word for a list) of \(N\) integers and must find a contiguous section of the array (often called a subarray, or an interval) that contains exactly \(K\) elements and has the largest sum.

One straightforward approach is to note that there are only \(N-K+1\) possible intervals (watch out for off-by-one errors!), so we could try summing up each interval and taking the best one. The code looks something like this:

• Python 3
• C++

inputFile = open("manin.txt", "r")
N, K = map(int, inputFile.readline().split())
 
people = []
 
for i in range(N):
    people.append(int(inputFile.readline()))
 
# Tracks the interval with the largest sum
# we've seen so far.
most_people = -1
 
for start in range(N - K + 1):
    # Find the sum of the interval from [start..start+K-1]
    current_people = 0
    for i in range(K):
        current_people += people[start+i]
 
    # If the current interval has a larger sum
    # than the most we've seen, then update.
    if current_people > most_people:
        most_people = current_people
 
outputFile = open("manout.txt", "w")
outputFile.write(f"{most_people}\n")

#include<fstream>

using namespace std;

ifstream in("manin.txt");
ofstream out("manout.txt");

const int MAX_N = 1e5+5;

int N, K;
int people[MAX_N];

int main() {
    in >> N >> K;

    for(int i = 0; i < N; i++) {
        in >> people[i];
    }

    // Tracks the interval with the largest sum
    // we've seen so far.
    int most_people = -1;

    for(int start = 0; start < N-K+1; start++) {
        // Find the sum of the interval from [start..start+K-1]
        int current_people = 0;
        for(int i = 0; i < K; i++) {
            current_people += people[start+i];
        }

        // If the current interval has a larger sum
        // than the most we've seen, then update.
        if(current_people > most_people) {
            most_people = current_people;
        }
    }

    out << most_people << "\n";
}

The worst-case time complexity of this approach is \(O(NK)\) (check for yourself that this is true) and would only score 50 out of 100 points for this problem. The part that's slow is the inner for-loop that calculates the sum of each interval as we need it. Is there a better way?

Prefix sums

We can start our algorithm by precalculating a prefix sum table (a.k.a cumulative sums, partial sums) that lets us calculate the sum of each interval much faster. Consider this example:

The blue (top) interval can be calculated by taking the red (bottom) interval and subtracting the green (middle) interval. Notice that the green and red intervals are both prefixes -- that is, they start from the beginning of the array. What's useful here, is that the sum of any interval of the array can be calculated as the difference between two prefixes.

This means that if we precalculate the sum of every prefix of an array, we can then calculate the sum of any interval (not just prefixes) in \(O(1)\) time! Calculating the prefix sums takes \(O(N)\) time if done efficiently. This allows us to get 100 points on Mansion.

In summary, prefix sums are a data structure that can be created in \(O(N)\) time and let you query the sum of any interval in \(O(1)\) time. It's a very powerful building block for more complex algorithms!

Exercises

Mansion (AIO 2007)

Using your newfound knowledge, try to solve Mansion.

Alien (AIO 2011)

Time to put your problem solving skills to the test! Go ahead and try to solve Alien from AIO 2011.

Hint if you're stuck (spend at least 15 minutes thinking about it yourself)

Click here to go back.