Code Implementation

class Solution:
    def matrixBlockSum(self, mat, K):
        m, n = len(mat), len(mat[0])
        # Compute prefix sum
        prefix = [[0]*(n+1) for _ in range(m+1)]
        for i in range(m):
            for j in range(n):
                prefix[i+1][j+1] = mat[i][j] + prefix[i][j+1] + prefix[i+1][j] - prefix[i][j]
        # Compute block sum for each cell
        result = [[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)
                result[i][j] = prefix[r2+1][c2+1] - prefix[r1][c2+1] - prefix[r2+1][c1] + prefix[r1][c1]
        return result
      

class Solution {
public:
    vector<vector<int>> matrixBlockSum(vector<vector<int>>& mat, int K) {
        int m = mat.size(), n = mat[0].size();
        vector<vector<int>> prefix(m+1, vector<int>(n+1, 0));
        for (int i = 0; i < m; ++i)
            for (int j = 0; j < n; ++j)
                prefix[i+1][j+1] = mat[i][j] + prefix[i][j+1] + prefix[i+1][j] - prefix[i][j];
        vector<vector<int>> result(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);
                result[i][j] = prefix[r2+1][c2+1] - prefix[r1][c2+1] - prefix[r2+1][c1] + prefix[r1][c1];
            }
        }
        return result;
    }
};
      

class Solution {
    public int[][] matrixBlockSum(int[][] mat, int K) {
        int m = mat.length, n = mat[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] = mat[i][j] + prefix[i][j+1] + prefix[i+1][j] - prefix[i][j];
        int[][] result = 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);
                result[i][j] = prefix[r2+1][c2+1] - prefix[r1][c2+1] - prefix[r2+1][c1] + prefix[r1][c1];
            }
        }
        return result;
    }
}
      

var matrixBlockSum = function(mat, K) {
    const m = mat.length, n = mat[0].length;
    const prefix = Array.from({length: m+1}, () => Array(n+1).fill(0));
    for (let i = 0; i < m; ++i)
        for (let j = 0; j < n; ++j)
            prefix[i+1][j+1] = mat[i][j] + prefix[i][j+1] + prefix[i+1][j] - prefix[i][j];
    const result = Array.from({length: m}, () => Array(n).fill(0));
    for (let i = 0; i < m; ++i) {
        for (let j = 0; j < n; ++j) {
            let r1 = Math.max(0, i-K), c1 = Math.max(0, j-K);
            let r2 = Math.min(m-1, i+K), c2 = Math.min(n-1, j+K);
            result[i][j] = prefix[r2+1][c2+1] - prefix[r1][c2+1] - prefix[r2+1][c1] + prefix[r1][c1];
        }
    }
    return result;
};
      

Problem Description

Given a matrix mat of size m x n and an integer K, your task is to build a new matrix answer where each answer[i][j] is equal to the sum of all elements mat[r][c] such that:

• i - K <= r <= i + K
• j - K <= c <= j + K
• and r and c are valid indices within the matrix bounds

In other words, for each cell, you sum all the values in a square block centered at that cell, with a "radius" of K, and taking care not to go outside the matrix.

Constraints:

• 1 <= m, n <= 100
• 0 <= mat[i][j] <= 100
• 0 <= K <= min(m, n)

Thought Process

At first glance, the problem seems straightforward: for each cell in the matrix, consider all cells within a square of side 2K+1 centered at that cell, and sum their values. This could be done with a nested loop for each cell, but that would be inefficient for large matrices, as it would require checking up to (2K+1)^2 cells for every cell in the matrix.

To optimize, we want a way to quickly compute the sum of any rectangular region in the matrix. This is a classic use-case for 2D prefix sums (also known as summed-area tables). By precomputing a prefix sum matrix, we can calculate the sum of any submatrix in constant time, reducing the overall complexity drastically.

Solution Approach

The solution uses a 2D prefix sum array to efficiently compute the block sum for each cell:

1. Build a Prefix Sum Matrix:
   • Create a matrix prefix where prefix[i+1][j+1] stores the sum of all elements in the rectangle from the top-left (0,0) to (i,j).
   • We use an extra row and column (of zeros) to handle boundary cases gracefully.
2. Calculate Block Sums:
   • For each cell (i, j) in the matrix, determine the boundaries of the square block using K (make sure not to go out of bounds).
   • Use the inclusion-exclusion principle with the prefix sum matrix to compute the sum of the block in constant time.
3. Return the Result:
   • Store the computed block sum for each cell in the result matrix.

The key design choice is using the prefix sum matrix, which allows for quick region sum lookups and avoids recomputation.

Example Walkthrough

Let's use the following sample input:

mat = [
  [1, 2, 3],
  [4, 5, 6],
  [7, 8, 9]
]
K = 1
  

For K = 1, the block sum for cell (1,1) (the center) includes all cells, so the sum is 45.

• For cell (0,0), the block covers cells (0,0), (0,1), (1,0), (1,1): 1+2+4+5=12.
• For cell (0,1), the block covers (0,0)-(1,2): 1+2+3+4+5+6=21.
• For cell (2,2), the block covers (1,1)-(2,2): 5+6+8+9=28.

The prefix sum matrix lets us compute each of these block sums in constant time, by using the inclusion-exclusion formula:

• For cell (i, j), block corners are (r1, c1) and (r2, c2), and block sum is:
  prefix[r2+1][c2+1] - prefix[r1][c2+1] - prefix[r2+1][c1] + prefix[r1][c1]

This process is repeated for every cell, efficiently filling the output matrix.

Time and Space Complexity

Brute-force Approach:

• For each cell, sum up to (2K+1)^2 cells.
• Total time: O(m * n * (2K+1)^2)
• Space: O(mn) for the result matrix.

• Building the prefix sum matrix: O(mn)
• For each cell, block sum lookup: O(1)
• Total time: O(mn)
• Space: O(mn) for both prefix and result matrices.

The optimized approach is much faster, especially for large values of K and large matrices.

Summary

The Matrix Block Sum problem is elegantly solved using a 2D prefix sum (summed-area table), which allows for constant-time queries of any rectangular region's sum. By precomputing the prefix sums, we avoid redundant calculations and achieve a highly efficient solution. This technique is a powerful pattern for a wide range of matrix region sum problems, and understanding it unlocks fast solutions to many related challenges.