Prefix Sum Applications: Subarray Sum Equals K & 2D Prefix Sums

Once you understand prefix sums, two powerful patterns open up: counting subarrays with a target sum in O(n) using a hash map, and answering any submatrix sum in O(1) with a 2D prefix sum. Both reuse the same “the answer is a difference of two prefixes” insight.

Subarray sum equals k

Count the number of contiguous subarrays whose elements sum to exactly k. Brute force is O(n²). The trick: a subarray a[l..r] sums to k exactly when prefix[r+1] - prefix[l] == k, i.e. prefix[l] == prefix[r+1] - k. So as we sweep, we keep a hash map of how many times each prefix sum has occurred, and for each new prefix we look up how many earlier prefixes equal current - k.

#include <vector>
#include <unordered_map>
int subarraySum(const std::vector<int>& a, int k) {
    std::unordered_map<long long, int> seen; // prefix value -> count
    seen[0] = 1;                              // empty prefix
    long long sum = 0;
    int count = 0;
    for (int x : a) {
        sum += x;
        auto it = seen.find(sum - k);
        if (it != seen.end()) count += it->second;
        seen[sum]++;
    }
    return count;
}

import java.util.HashMap;
import java.util.Map;
int subarraySum(int[] a, int k) {
    Map<Long, Integer> seen = new HashMap<>(); // prefix value -> count
    seen.put(0L, 1);                           // empty prefix
    long sum = 0;
    int count = 0;
    for (int x : a) {
        sum += x;
        count += seen.getOrDefault(sum - k, 0);
        seen.merge(sum, 1, Integer::sum);
    }
    return count;
}

function subarraySum(a, k) {
  const seen = new Map(); // prefix value -> count
  seen.set(0, 1);         // empty prefix
  let sum = 0, count = 0;
  for (const x of a) {
    sum += x;
    count += seen.get(sum - k) || 0;
    seen.set(sum, (seen.get(sum) || 0) + 1);
  }
  return count;
}

def subarray_sum(a, k):
    seen = {0: 1}   # prefix value -> count, empty prefix
    s = 0
    count = 0
    for x in a:
        s += x
        count += seen.get(s - k, 0)
        seen[s] = seen.get(s, 0) + 1
    return count

For a = [1, 1, 1], k = 2, the answer is 2 ([1,1] at indices 0-1 and 1-2). O(n) time, O(n) space.

For beginners: Seeding the map with {0: 1} represents the empty prefix — it’s what makes subarrays that start at index 0 count correctly. Forgetting it is the single most common bug here.

Going deeper: This hash-map-of-prefixes idea generalizes far beyond sums: prefix-XOR counts subarrays with a target XOR, and prefix-sum-mod-k counts subarrays divisible by k. Same skeleton, different key.

2D prefix sums (submatrix sum in O(1))

For a grid, build a 2D prefix array P where P[i][j] is the sum of the rectangle from the top-left corner (0,0) to (i-1, j-1). With a one-row and one-column border of zeros, every cell is:

P[i][j] = grid[i-1][j-1] + P[i-1][j] + P[i][j-1] - P[i-1][j-1]

We subtract P[i-1][j-1] because the overlapping top-left block got added twice (inclusion-exclusion). Then the sum of the submatrix with rows r1..r2 and columns c1..c2 (inclusive) is:

sum = P[r2+1][c2+1] - P[r1][c2+1] - P[r2+1][c1] + P[r1][c1]

#include <vector>
struct Matrix2D {
    std::vector<std::vector<long long>> P;
    Matrix2D(const std::vector<std::vector<int>>& g) {
        int m = g.size(), n = g[0].size();
        P.assign(m + 1, std::vector<long long>(n + 1, 0));
        for (int i = 0; i < m; i++)
            for (int j = 0; j < n; j++)
                P[i+1][j+1] = g[i][j] + P[i][j+1] + P[i+1][j] - P[i][j];
    }
    long long query(int r1, int c1, int r2, int c2) { // inclusive
        return P[r2+1][c2+1] - P[r1][c2+1] - P[r2+1][c1] + P[r1][c1];
    }
};

class Matrix2D {
    long[][] P;
    Matrix2D(int[][] g) {
        int m = g.length, n = g[0].length;
        P = new long[m + 1][n + 1];
        for (int i = 0; i < m; i++)
            for (int j = 0; j < n; j++)
                P[i+1][j+1] = g[i][j] + P[i][j+1] + P[i+1][j] - P[i][j];
    }
    long query(int r1, int c1, int r2, int c2) { // inclusive
        return P[r2+1][c2+1] - P[r1][c2+1] - P[r2+1][c1] + P[r1][c1];
    }
}

class Matrix2D {
  constructor(g) {
    const m = g.length, n = g[0].length;
    this.P = Array.from({ length: m + 1 }, () => new Array(n + 1).fill(0));
    for (let i = 0; i < m; i++)
      for (let j = 0; j < n; j++)
        this.P[i+1][j+1] = g[i][j] + this.P[i][j+1] + this.P[i+1][j] - this.P[i][j];
  }
  query(r1, c1, r2, c2) { // inclusive
    const P = this.P;
    return P[r2+1][c2+1] - P[r1][c2+1] - P[r2+1][c1] + P[r1][c1];
  }
}

class Matrix2D:
    def __init__(self, g):
        m, n = len(g), len(g[0])
        self.P = [[0] * (n + 1) for _ in range(m + 1)]
        for i in range(m):
            for j in range(n):
                self.P[i+1][j+1] = (g[i][j] + self.P[i][j+1]
                                    + self.P[i+1][j] - self.P[i][j])

    def query(self, r1, c1, r2, c2):  # inclusive
        P = self.P
        return P[r2+1][c2+1] - P[r1][c2+1] - P[r2+1][c1] + P[r1][c1]

Build: O(m·n) time and space.
Each submatrix query: O(1).

For beginners: The four-term query formula is just inclusion-exclusion: take the big rectangle, subtract the strip above and the strip to the left, then add back the top-left corner you subtracted twice. Draw it once on paper and it sticks.

Common pitfalls

Skipping the {0: 1} seed in subarray-sum-equals-k — subarrays starting at index 0 get miscounted.
Updating the map before the lookup. Look up sum - k first, then record sum, or a zero-length match can sneak in.
2D border mistakes. The padded prefix array is (m+1) x (n+1); query indices are shifted by one. Mixing inclusive/exclusive bounds is the usual source of bugs.
Overflow on large grids or values — use 64-bit accumulators.

Practice these on the exercises, or revisit the foundation in prefix sums.