Prefix Sums

cs 26-12-2023

Merry christmas!

Imagine you are managing an inventory of items where a certain number of items will get added to or removed from that inventory over a period of time. You are provided with a log of the number of items that are added to the inventory at each time unit, and you want to find the net change in the number of items over a given period of time.

Naively, you can do this by simply summing up the changes over the period of time specified in the query. Since, for a log that is n entries long, the worst case scenario would be summing up the changes over all N entries, and thus the time complexity is O(n).

For a few smaller queries, this is acceptable. However, once you have thousands upon thousands of items removed or added, just like in a real inventory, as well as thousands of queries coming in every second, this is no longer sufficient. Luckily, there is a solution:

Prefix sums allows the same answer to be reached in a time complexity of just O(1). Instead of having to do the same summing calculation for each query, we can first calculate a cumulative sum over the entire log so that for each query, we can simply subtract the first and last elements corresponding to our time range.

The following example should illustrate this more concretely:

USACO 2016 January Problem 2: http://www.usaco.org/index.php?page=viewproblem2&cpid=595

If we were to do it with brute force, the time complexity would be O(n^3): O(n^2) for a nested for loop covering all sets of cows, and another O(n) for another loop summing up the IDs of each cow within each set. However, since O(n^3) is unacceptable, we need to look for other methods.

Fortunately, we realise that prefix sums allow us to skip the additional O(n) loop at the end for summing up the IDs—reducing our time complexity to O(n^2) only. This is sufficient for our purposes as N < 50,000.

int n;
cin >> n;
vector<long long> sums(n + 1);
for (int i = 0; i < n; i++) {
    sums[i + 1] = sums[i];
    long long id;
    cin >> id;
    sums[i + 1] += id;
}

This code shows the formation of our prefix sum vector, called sums. Each iteration, we copy the previous value in sums to the current index, then add the inputted id to that value to form our cumulative sum.

Finally, we can take the prefix sums and iterate through every subset, outputting the maximum of each difference as our answer. The final code is shown below:

#include <iostream>
#include <vector>

using namespace std;

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    freopen("div7.in", "r", stdin);
    freopen("div7.out", "w", stdout);

    int n;
    cin >> n;
    vector<long long> sums(n + 1);
    for (int i = 0; i < n; i++) {
        sums[i + 1] = sums[i];
        long long id;
        cin >> id;
        sums[i + 1] += id;
    }

    int max_group = 0;
    for (int i = 0; i < n; i++) {
        for (int j = n; j > i; j--) {
            if ((sums[j] - sums[i]) % 7 == 0) {
                max_group = max(max_group, j - i);
            }
        }
    }

    cout << max_group << "\n";
}
Author's photo

Lin Jiang

I’m a high school student that's passionate about computer science. Join me in having fun through recreational programming!

See other articles: