1314. Matrix Block Sum
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.length
n == mat[i].length
1 <= m, n, k <= 100
1 <= mat[i][j] <= 100
Solution
matrix-block-sum.py
class Solution:
def matrixBlockSum(self, mat: List[List[int]], k: int) -> List[List[int]]:
rows, cols = len(mat), len(mat[0])
prefix = [[0] * (cols + 1) for _ in range(rows + 1)]
for i in range(rows):
for j in range(cols):
prefix[i + 1][j + 1] = prefix[i + 1][j] + prefix[i][j + 1] - prefix[i][j] + mat[i][j]
for i in range(rows):
for j in range(cols):
r1, r2 = max(0, i - k), min(rows, i + k + 1)
c1, c2 = max(0, j - k), min(cols, j + k + 1)
mat[i][j] = (prefix[r2][c2] - prefix[r1][c2]) - (prefix[r2][c1] - prefix[r1][c1])
return mat