Problem

Given a m x n matrix mat and an integer k, return a matrix answer where each answer[i][j] is the sum of all elements mat[r][c] for:

  • i - k <= r <= i + k,
  • j - k <= c <= j + k, and
  • (r, c) is a valid position in the matrix.

Examples

Example 1:

1
2
3
4

Input:
mat = [[1,2,3],[4,5,6],[7,8,9]], k = 1
Output:
 [[12,21,16],[27,45,33],[24,39,28]]

Example 2:

1
2
3
4

Input:
mat = [[1,2,3],[4,5,6],[7,8,9]], k = 2
Output:
 [[45,45,45],[45,45,45],[45,45,45]]

Solution

Method 1 - Using Prefix sum array

Here is the approach:

  1. Use a prefix sum array to efficiently calculate the sum of elements within any submatrix. The prefix sum array P[i][j] will store the sum of elements in the submatrix starting from (0, 0) to (i, j).
  2. For each element in the answer matrix, use the prefix sum array to calculate the sum of elements within the given distance k.

Code

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29

public class Solution {
    public int[][] matrixBlockSum(int[][] mat, int k) {
        int m = mat.length;
        int n = mat[0].length;
        int[][] answer = new int[m][n];
        int[][] prefixSum = new int[m + 1][n + 1];

        // Compute the prefix sum array
        for (int i = 1; i <= m; i++) {
            for (int j = 1; j <= n; j++) {
                prefixSum[i][j] = mat[i - 1][j - 1] + prefixSum[i - 1][j] + prefixSum[i][j - 1] - prefixSum[i - 1][j - 1];
            }
        }

        // Compute the result matrix
        for (int i = 0; i < m; i++) {
            for (int j = 0; j < n; j++) {
                int r1 = Math.max(0, i - k);
                int c1 = Math.max(0, j - k);
                int r2 = Math.min(m - 1, i + k);
                int c2 = Math.min(n - 1, j + k);

                answer[i][j] = prefixSum[r2 + 1][c2 + 1] - prefixSum[r1][c2 + 1] - prefixSum[r2 + 1][c1] + prefixSum[r1][c1];
            }
        }

        return answer;
    }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22

class Solution:
    def matrixBlockSum(self, mat: List[List[int]], k: int) -> List[List[int]]:
        m, n = len(mat), len(mat[0])
        answer = [[0] * n for _ in range(m)]
        prefix_sum = [[0] * (n + 1) for _ in range(m + 1)]

        # Compute the prefix sum array
        for i in range(1, m + 1):
            for j in range(1, n + 1):
                prefix_sum[i][j] = mat[i-1][j-1] + prefix_sum[i-1][j] + prefix_sum[i][j-1] - prefix_sum[i-1][j-1]

        # Compute the result matrix
        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)

                answer[i][j] = prefix_sum[r2+1][c2+1] - prefix_sum[r1][c2+1] - prefix_sum[r2+1][c1] + prefix_sum[r1][c1]

        return answer

Complexity

  • ⏰ Time complexity: O(m * n), where m is the number of rows and n is the number of columns, as we preprocess the matrix to create the prefix sum array and then compute the result matrix.
  • 🧺 Space complexity: O(m * n), for storing the prefix sum array and the result matrix.