Description
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.
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.lengthn == mat[i].length1 <= m, n, k <= 1001 <= mat[i][j] <= 100
Solutions
Solution 1: Two-Dimensional Prefix Sum
This problem is a template for two-dimensional prefix sum.
We define s[i][j] as the sum of the elements in the first i rows and the first j columns of the matrix mat. The calculation formula for s[i][j] is:
s[i][j] = s[i-1][j] + s[i][j-1] - s[i-1][j-1] + mat[i-1][j-1]
In this way, we can quickly calculate the sum of elements in any rectangular area through the s array.
For a rectangular area with the upper left coordinate (x1, y1) and the lower right coordinate (x2, y2), we can calculate the sum of its elements through the s array:
s[x2+1][y2+1] - s[x1][y2+1] - s[x2+1][y1] + s[x1][y1]
The time complexity is O(m × n), and the space complexity is O(m × n). Where m and n are the number of rows and columns in the matrix, respectively.
PythonJavaC++GoTypeScript
class Solution:
def matrixBlockSum(self, mat: List[List[int]], k: int) -> List[List[int]]:
m, n = len(mat), len(mat[0])
s = [[0] * (n + 1) for _ in range(m + 1)]
for i, row in enumerate(mat, 1):
for j, x in enumerate(row, 1):
s[i][j] = s[i - 1][j] + s[i][j - 1] - s[i - 1][j - 1] + x
ans = [[0] * n for _ in range(m)]
for i in range(m):
for j in range(n):
x1, y1 = max(i - k, 0), max(j - k, 0)
x2, y2 = min(m - 1, i + k), min(n - 1, j + k)
ans[i][j] = (
s[x2 + 1][y2 + 1] - s[x1][y2 + 1] - s[x2 + 1][y1] + s[x1][y1]
)
return ans(code-box)
class Solution {
public int[][] matrixBlockSum(int[][] mat, int k) {
int m = mat.length;
int n = mat[0].length;
int[][] s = new int[m + 1][n + 1];
for (int i = 0; i < m; ++i) {
for (int j = 0; j < n; ++j) {
s[i + 1][j + 1] = s[i][j + 1] + s[i + 1][j] - s[i][j] + mat[i][j];
}
}
int[][] ans = new int[m][n];
for (int i = 0; i < m; ++i) {
for (int j = 0; j < n; ++j) {
int x1 = Math.max(i - k, 0);
int y1 = Math.max(j - k, 0);
int x2 = Math.min(m - 1, i + k);
int y2 = Math.min(n - 1, j + k);
ans[i][j] = s[x2 + 1][y2 + 1] - s[x1][y2 + 1] - s[x2 + 1][y1] + s[x1][y1];
}
}
return ans;
}
}(code-box)
class Solution {
public:
vector<vector<int>> matrixBlockSum(vector<vector<int>>& mat, int k) {
int m = mat.size();
int n = mat[0].size();
vector<vector<int>> s(m + 1, vector<int>(n + 1));
for (int i = 0; i < m; ++i) {
for (int j = 0; j < n; ++j) {
s[i + 1][j + 1] = s[i][j + 1] + s[i + 1][j] - s[i][j] + mat[i][j];
}
}
vector<vector<int>> ans(m, vector<int>(n));
for (int i = 0; i < m; ++i) {
for (int j = 0; j < n; ++j) {
int x1 = max(i - k, 0);
int y1 = max(j - k, 0);
int x2 = min(m - 1, i + k);
int y2 = min(n - 1, j + k);
ans[i][j] = s[x2 + 1][y2 + 1] - s[x1][y2 + 1] - s[x2 + 1][y1] + s[x1][y1];
}
}
return ans;
}
};(code-box)
func matrixBlockSum(mat [][]int, k int) [][]int {
m, n := len(mat), len(mat[0])
s := make([][]int, m+1)
for i := range s {
s[i] = make([]int, n+1)
}
for i, row := range mat {
for j, x := range row {
s[i+1][j+1] = s[i][j+1] + s[i+1][j] - s[i][j] + x
}
}
ans := make([][]int, m)
for i := range ans {
ans[i] = make([]int, n)
}
for i := 0; i < m; i++ {
for j := 0; j < n; j++ {
x1 := max(i-k, 0)
y1 := max(j-k, 0)
x2 := min(m-1, i+k)
y2 := min(n-1, j+k)
ans[i][j] = s[x2+1][y2+1] - s[x1][y2+1] - s[x2+1][y1] + s[x1][y1]
}
}
return ans
}(code-box)
function matrixBlockSum(mat: number[][], k: number): number[][] {
const m: number = mat.length;
const n: number = mat[0].length;
const s: number[][] = Array.from({ length: m + 1 }, () => Array(n + 1).fill(0));
for (let i = 0; i < m; i++) {
for (let j = 0; j < n; j++) {
s[i + 1][j + 1] = s[i][j + 1] + s[i + 1][j] - s[i][j] + mat[i][j];
}
}
const ans: number[][] = Array.from({ length: m }, () => Array(n).fill(0));
for (let i = 0; i < m; i++) {
for (let j = 0; j < n; j++) {
const x1: number = Math.max(i - k, 0);
const y1: number = Math.max(j - k, 0);
const x2: number = Math.min(m - 1, i + k);
const y2: number = Math.min(n - 1, j + k);
ans[i][j] = s[x2 + 1][y2 + 1] - s[x1][y2 + 1] - s[x2 + 1][y1] + s[x1][y1];
}
}
return ans;
}(code-box)
