DDSA Solutions

2D Submatrix Sum Queries

Advertisement

Intuition

Answer submatrix sum queries in O(1) using 2D prefix sums.

Algorithm

  1. 1prefix[i][j] = sum of matrix[0..i-1][0..j-1]. Query(r1,c1,r2,c2) = prefix[r2+1][c2+1] - prefix[r1][c2+1] - prefix[r2+1][c1] + prefix[r1][c1].

Common Pitfalls

  • Inclusion-exclusion principle. Build prefix in O(m*n), each query O(1). Same as LC 304.
2D Submatrix Sum Queries.java
Java
// Approach: 2D prefix sum (integral image). Precompute prefix[i][j] = sum of all cells (0,0) to (i-1,j-1).
// Answer any submatrix query in O(1) using inclusion-exclusion on prefix sums.
// Time: O(n*m) build, O(1) query Space: O(n*m)
import java.util.*;

class Solution {
    public ArrayList<Integer> prefixSum2D(int[][] mat, int[][] queries) {
        int n = mat.length;
        int m = mat[0].length;

        // Step 1: Build prefix sum matrix
        int[][] ps = new int[n][m];

        for (int i = 0; i < n; i++) {
            for (int j = 0; j < m; j++) {
                ps[i][j] = mat[i][j];

                if (i > 0)
                    ps[i][j] += ps[i - 1][j];
                if (j > 0)
                    ps[i][j] += ps[i][j - 1];
                if (i > 0 && j > 0)
                    ps[i][j] -= ps[i - 1][j - 1];
            }
        }

        // Step 2: Process queries
        ArrayList<Integer> result = new ArrayList<>();

        for (int[] q : queries) {
            int r1 = q[0], c1 = q[1], r2 = q[2], c2 = q[3];

            int sum = ps[r2][c2];

            if (r1 > 0)
                sum -= ps[r1 - 1][c2];
            if (c1 > 0)
                sum -= ps[r2][c1 - 1];
            if (r1 > 0 && c1 > 0)
                sum += ps[r1 - 1][c1 - 1];

            result.add(sum);
        }

        return result;
    }
}
Advertisement
Was this solution helpful?