Prefix sums

The pattern

Build an array where prefix[i] is the sum of the first i elements. Then the sum of any range is one subtraction.

sum(i .. j) = prefix[j+1] - prefix[i]

You pay O(n) once, and every query after that is O(1). If a problem asks the same kind of question many times over a fixed array, precomputing is almost always the intended answer.

Recognise it from: sum of a range, many queries, subarray sum, average of a segment, equilibrium point.

Building it

Use a prefix array of size n+1 with prefix[0] = 0. The extra slot removes every special case for ranges starting at index 0.

#include <iostream>
#include <vector>
#include <unordered_map>
using namespace std;

int main() {
    vector<int> a = {3, 1, 4, 1, 5, 9, 2, 6};
    int n = a.size();

    vector<long long> prefix(n + 1, 0);
    for (int i = 0; i < n; i++) prefix[i + 1] = prefix[i] + a[i];

    // prefix = 0 3 4 8 9 14 23 25 31

    // sum of a[2..5] = 4 + 1 + 5 + 9 = 19
    cout << prefix[6] - prefix[2] << "\n";     // 19

    // sum of the whole array
    cout << prefix[n] - prefix[0] << "\n";     // 31
    return 0;
}

Without prefix sums, answering q queries costs O(n) each, so O(n x q). With them it is O(n + q). For n = 100,000 and q = 100,000 that is 10^10 versus 200,000.

The problem

Count the number of contiguous subarrays whose sum equals exactly k.

nums = {1, 2, 3, -3, 1, 1, 1}, k = 3

Step 1: brute force

int count_brute(const vector<int> &a, int k) {
    int n = a.size(), count = 0;
    for (int i = 0; i < n; i++) {
        long long sum = 0;
        for (int j = i; j < n; j++) {
            sum += a[j];
            if (sum == k) count++;
        }
    }
    return count;
}

O(n^2). It passes for n = 1,000 and fails for n = 100,000.

Step 2: rewrite the condition

sum(i .. j) = prefix[j+1] - prefix[i]

We want that to equal k:

prefix[j+1] - prefix[i] = k, so prefix[i] = prefix[j+1] - k

So at each position j, the question becomes: how many earlier prefix values equal (current prefix - k)? Counting "how many times have I seen this value" is exactly what a hash map does, in O(1).

The O(n^2) search over pairs became an O(1) lookup. That combination - prefix sums plus a hash map - is one of the most reused ideas in interviews.

Step 3: the code

int subarray_sum_equals_k(const vector<int> &a, int k) {
    unordered_map<long long, int> seen;
    seen[0] = 1;                 // the empty prefix, seen once

    long long running = 0;
    int count = 0;

    for (int x : a) {
        running += x;
        auto it = seen.find(running - k);
        if (it != seen.end()) count += it->second;
        seen[running]++;
    }
    return count;
}

int main() {
    vector<int> a = {1, 2, 3, -3, 1, 1, 1};
    cout << subarray_sum_equals_k(a, 3) << "\n";   // 6
    vector<int> b = {1, 1, 1};
    cout << subarray_sum_equals_k(b, 2) << "\n";   // 2
    return 0;
}

O(n) time, O(n) space.

For {1, 2, 3, -3, 1, 1, 1} with k = 3 the six subarrays are {1,2}, {1,2,3,-3}, {2,3,-3,1}, {3}, {3,-3,1,1,1} and {1,1,1}. The prefix array is 0 1 3 6 3 4 5 6, and every pair of equal-difference-3 entries in it is one of those six. Count them by hand once - it is the fastest way to believe the method.

The 2D version

The same idea in two dimensions gives the sum of any rectangle in O(1).

// pre[i][j] = sum of the rectangle from (0,0) to (i-1,j-1)
// sum of rows r1..r2 and cols c1..c2:
//   pre[r2+1][c2+1] - pre[r1][c2+1] - pre[r2+1][c1] + pre[r1][c1]

That last + pre[r1][c1] is inclusion-exclusion: the top-left block was subtracted twice, so add it back once. Draw the four rectangles on paper once and you will never forget it.

Practice, in this order

1. Running sum of an array - the one-line warm-up.
2. Find the pivot index, where left sum equals right sum.
3. Contiguous array - longest subarray with equal 0s and 1s. Map 0 to -1 and

it becomes "longest subarray with sum 0".

1. Subarray sums divisible by k, storing running % k instead of running.

Be careful: C++ gives negative remainders, so use ((running % k) + k) % k.