The O(n) Club: The 2D Prefix Sum – Stop Hating Matrix Math (and Yourself)

The O(n) Club: The 2D Prefix Sum – Stop Hating Matrix Math (and Yourself)

⚡ TL;DR

Quit writing those nested loops for every sum in a matrix. Precompute a 2D prefix sum array that’s just one row and column bigger—then get any rectangle’s sum instantly with one slightly magical formula. Yes, it’s that easy. Here’s how simple it looks in Java:

// Precompute in O(mn)
int[][] prefix = ... // see below
// Query sum for (row1,col1)-(row2,col2) in O(1)
int total = prefix[row2+1][col2+1] - prefix[row2+1][col1] - prefix[row1][col2+1] + prefix[row1][col1];

🧠 How Devs Usually Mess This Up

Step 1: Write naive code that works for a 3×3 grid, then watch it crawl for any interview-sized input. Step 2: Read one blog post, try to direct-move prefix sums without that extra row/column. Prepare your resignation letter after wrestling with off-by-one madness. Step 3: Flip exactly one sign or index and stare blankly as every test case fails in new and innovative ways. Bonus round: Try updating the matrix dynamically (spoiler—different data structure, different flavor of suffering).

🔍 What’s Actually Going On

If you want to sum values inside any rectangle of your 2D array fast, you need a little forward thinking: the 2D prefix sum array (a.k.a. summed-area table, if you enjoy making people yawn). Make a new array, one bigger on each side, where s[i+1][j+1] is the sum for the rectangle from (0,0) to (i,j). Now, with careful inclusion-exclusion (subtract above and left, then add back the overlap), you can get area sums in O(1) time. That extra row and column? They exist so you never have to write if (i > 0) again. You’re welcome.

🛠️ PseudoCode

• Precompute the prefix sum grid:

  // Java: for matrix[m][n]
  int m = matrix.length, n = matrix[0].length;
  int[][] prefix = new int[m+1][n+1];
  for (int i = 0; i < m; i++) {
    for (int j = 0; j < n; j++) {
      prefix[i+1][j+1] = matrix[i][j]
                        + prefix[i][j+1]
                        + prefix[i+1][j]
                        - prefix[i][j];
    }
  }
  

  • Adds current cell
  • Adds area above
  • Adds area to the left
  • Subtracts the top-left overlap (yeah, you double-counted it)
• Query any submatrix (row1, col1) to (row2, col2) inclusive:

  int sum = prefix[row2+1][col2+1]
          - prefix[row2+1][col1]
          - prefix[row1][col2+1]
          + prefix[row1][col1];

  • Take the big area up to bottom-right
  • Subtract area above
  • Subtract area to the left
  • Add back overlap (you were too eager with the minus)

💻 The Code

class NumMatrix {
    private int[][] prefix;
     public NumMatrix(int[][] matrix) {
        int m = matrix.length, n = matrix[0].length;
        prefix = new int[m + 1][n + 1];
        for (int i = 0; i < m; i++) {
            for (int j = 0; j < n; j++) {
                prefix[i + 1][j + 1] = matrix[i][j]
                                      + prefix[i][j + 1]
                                      + prefix[i + 1][j]
                                      - prefix[i][j];
            }
        }
    }
     // Returns the sum for (row1,col1) to (row2,col2), inclusive
    public int sumRegion(int row1, int col1, int row2, int col2) {
        return prefix[row2 + 1][col2 + 1]
             - prefix[row2 + 1][col1]
             - prefix[row1][col2 + 1]
             + prefix[row1][col1];
    }
}

⚠️ Pitfalls, Traps & Runtime Smacks

• Off-by-one errors: If your prefix size isn’t m+1 x n+1, break out the Tylenol now.
• Sign mistakes: A single plus or minus botched and your output will look like performance art. Triple-check the formula.
• Not for updates: This is for immutable matrices. If you need updates, go read up on Fenwick Trees (and pour a stiffer coffee).
• Extra space: Technically uses a whole extra array. But if you’re still worried about *that* in year 2024, you’re reading the wrong blog.
• Preprocess O(mn), Query O(1): If you ever see O(n²) in your interviewer’s eye, this is your lightsaber.

🧠 Interviewer Brain-Tattoo

If you can’t do 2D prefix sums, you’re one Google away from brute-force shame. Get that m+1/n+1 array right, and you’ll never fear region queries again.