Given a m * 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.
Example 1:
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:
Input: mat = [[1,2,3],[4,5,6],[7,8,9]], K = 2
Output: [[45,45,45],[45,45,45],[45,45,45]]
Constraints:
m == mat.length
n == mat[i].length
1 <= m, n, K <= 100
1 <= mat[i][j] <= 100
class Solution {
public int[][] matrixBlockSum(int[][] mat, int K) {
int m = mat.length, n = mat[0].length;
int prefixSum[][] = new int[m][n];
for (int i = 0; i < m; i++) {
int sum = 0;
for (int j = 0; j < n; j++) {
sum += mat[i][j];
if (i == 0)
prefixSum[i][j] = sum;
else
prefixSum[i][j] = sum + prefixSum[i - 1][j];
}
}
int[][] ans = new int[m][n];
for (int i = 0; i < m; i++) {
for (int j = 0; j < n; j++) {
int rUP = i + K, cUP = j + K;
if (rUP >= m)
rUP = m - 1;
if (cUP >= n)
cUP = n - 1;
ans[i][j] = prefixSum[rUP][cUP];
int r = i - K - 1;
int c = j - K - 1;
if (r < 0 && c < 0)
continue;
else if (r < 0)
ans[i][j] = ans[i][j] - prefixSum[rUP][c];
else if (c < 0)
ans[i][j] = ans[i][j] - prefixSum[r][cUP];
else
ans[i][j] = ans[i][j] - prefixSum[rUP][c] - prefixSum[r][cUP] + prefixSum[r][c];
}
}
return ans;
}
}