给你一个 m x n 的矩阵 mat 和一个整数 k ,请你返回一个矩阵 answer ,其中每个 answer[i][j] 是所有满足下述条件的元素 ma
t[r][c] 的和:
i - k <= r <= i + k,
j - k <= c <= j + k 且
(r, c) 在矩阵内。
示例 1:
输入:mat = [[1,2,3],[4,5,6],[7,8,9]], k = 1
输出:[[12,21,16],[27,45,33],[24,39,28]]
示例 2:
输入:mat = [[1,2,3],[4,5,6],[7,8,9]], k = 2
输出:[[45,45,45],[45,45,45],[45,45,45]]
提示:
m == mat.length
n == mat[i].length
1 <= m, n, k <= 100
1 <= mat[i][j] <= 100
Related Topics 数组 矩阵 前缀和
思路:
方法:二维前缀和
answer[i][j]的值是以mat[i][j]为中心的边长为 2k+1 的矩阵的和。
对原矩阵拓宽四周边界 k行 或者 k列,计算二维矩阵前缀和,使用标准的前缀和模板。
code:
import numpy as np
class Solution:
def matrixBlockSum(self, mat, k: int):
m, n = len(mat), len(mat[0])
# 拓宽矩阵, 计算二维前缀和
presum = [[0]*(n+2*k+1) for _ in range(0, m+2*k+1)]
# print(presum)
for i in range(k+1, m+2*k+1):
for j in range(k+1, n+2*k+1):
presum[i][j] = presum[i-1][j] + presum[i][j-1] - presum[i-1][j-1]
if (i-k-1) < m and j - k - 1 < n:
presum[i][j] += mat[i - k - 1][j - k -1]
print(presum)
# 矩阵和
answer = np.zeros((m, n), dtype=int).tolist()
print(answer)
for i in range(k+1, m+k+1):
for j in range(k+1, n+k+1):
answer[i-k-1][j-k-1] = presum[i+k][j+k] - presum[i-k-1][j+k] - presum[i+k][j-k-1] + presum[i-k-1][j-k-1]
return answer
s = Solution()
# print(s.matrixBlockSum(mat = [[1,2,3],[4,5,6],[7,8,9]], k = 1))
# print(s.matrixBlockSum(mat = [[1,2,3],[4,5,6],[7,8,9]], k = 2))
print(s.matrixBlockSum(mat=[[67,64,78],[99,98,38],[82,46,46],[6,52,55],[55,99,45]], k=3))