Matrix Block Sum

medium prefix sum matrix

Problem

Given an m × n matrix mat and an integer k, return a matrix answer where answer[i][j] equals the sum of all elements mat[r][c] with i − k ≤ r ≤ i + k and j − k ≤ c ≤ j + k, clamped to the grid.

Inputmat = [[1,2,3],[4,5,6],[7,8,9]], k = 1

Output[[12,21,16],[27,45,33],[24,39,28]]

answer[0][0] sums mat[0..1][0..1] = 1+2+4+5 = 12.

mat (JSON matrix)

def matrix_block_sum(mat, k):
    m, n = len(mat), len(mat[0])
    P = [[0] * (n + 1) for _ in range(m + 1)]
    for i in range(m):
        for j in range(n):
            P[i + 1][j + 1] = mat[i][j] + P[i][j + 1] + P[i + 1][j] - P[i][j]
    ans = [[0] * n for _ in range(m)]
    for i in range(m):
        for j in range(n):
            r1 = max(0, i - k); c1 = max(0, j - k)
            r2 = min(m - 1, i + k); c2 = min(n - 1, j + k)
            ans[i][j] = P[r2 + 1][c2 + 1] - P[r1][c2 + 1] - P[r2 + 1][c1] + P[r1][c1]
    return ans

function matrixBlockSum(mat, k) {
  const m = mat.length, n = mat[0].length;
  const 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++)
      P[i + 1][j + 1] = mat[i][j] + P[i][j + 1] + P[i + 1][j] - P[i][j];
  const ans = Array.from({ length: m }, () => new Array(n).fill(0));
  for (let i = 0; i < m; i++)
    for (let j = 0; j < n; j++) {
      const r1 = Math.max(0, i - k), c1 = Math.max(0, j - k);
      const r2 = Math.min(m - 1, i + k), c2 = Math.min(n - 1, j + k);
      ans[i][j] = P[r2 + 1][c2 + 1] - P[r1][c2 + 1] - P[r2 + 1][c1] + P[r1][c1];
    }
  return ans;
}

class Solution {
    public int[][] matrixBlockSum(int[][] mat, int k) {
        int m = mat.length, n = mat[0].length;
        int[][] P = new int[m + 1][n + 1];
        for (int i = 0; i < m; i++)
            for (int j = 0; j < n; j++)
                P[i + 1][j + 1] = mat[i][j] + P[i][j + 1] + P[i + 1][j] - P[i][j];
        int[][] ans = new int[m][n];
        for (int i = 0; i < m; i++) {
            for (int j = 0; j < n; j++) {
                int r1 = Math.max(0, i - k), c1 = Math.max(0, j - k);
                int r2 = Math.min(m - 1, i + k), c2 = Math.min(n - 1, j + k);
                ans[i][j] = P[r2 + 1][c2 + 1] - P[r1][c2 + 1] - P[r2 + 1][c1] + P[r1][c1];
            }
        }
        return ans;
    }
}

vector<vector<int>> matrixBlockSum(vector<vector<int>>& mat, int k) {
    int m = mat.size(), n = mat[0].size();
    vector<vector<int>> P(m + 1, vector<int>(n + 1, 0));
    for (int i = 0; i < m; i++)
        for (int j = 0; j < n; j++)
            P[i + 1][j + 1] = mat[i][j] + P[i][j + 1] + P[i + 1][j] - P[i][j];
    vector<vector<int>> ans(m, vector<int>(n, 0));
    for (int i = 0; i < m; i++) {
        for (int j = 0; j < n; j++) {
            int r1 = max(0, i - k), c1 = max(0, j - k);
            int r2 = min(m - 1, i + k), c2 = min(n - 1, j + k);
            ans[i][j] = P[r2 + 1][c2 + 1] - P[r1][c2 + 1] - P[r2 + 1][c1] + P[r1][c1];
        }
    }
    return ans;
}

Explanation

The naive way would re-add every cell inside each block, which is slow. Instead we build a 2D prefix-sum table P once, where P[r+1][c+1] holds the sum of the whole rectangle from the top-left corner down to row r, column c. After that, any block sum is just a quick formula.

Each prefix cell is filled with mat[i][j] + P[i][j+1] + P[i+1][j] - P[i][j]. We add the sums above and to the left, then subtract the top-left overlap that got counted twice. This is inclusion-exclusion.

For each output cell we clamp the block to the grid edges, giving corners (r1,c1) to (r2,c2). The block sum is P[r2+1][c2+1] - P[r1][c2+1] - P[r2+1][c1] + P[r1][c1] — the big rectangle minus the strip above, minus the strip to the left, plus the corner that was subtracted twice.

Example: for mat = [[1,2,3],[4,5,6],[7,8,9]] with k = 1, cell answer[0][0] covers rows 0..1 and columns 0..1, so it sums 1+2+4+5 = 12. Every cell is computed this way in constant time.

Time: O(m · n) Space: O(m · n)