题目描述:
给你一个 m * n 的矩阵 mat 和一个整数 K ,请你返回一个矩阵 answer ,其中每个 answer[i][j] 是所有满足下述条件的元素 mat[r][c] 的和:
i - K <= r <= i + K, j - K <= c <= j + K
(r, c) 在矩阵内。
思路:矩阵前缀和
首先求出矩阵的前缀和
然后再根据题目要求,求出子数组的元素和
i - K <= r <= i + K, j - K <= c <= j + K
也就是说以(i - K,j - K )为左上角,以( i + K,j + K)为右下角中间的所有子数组的和
代码如下:
C++:
class Solution {
public:
vector<vector<int>> matrixBlockSum(vector<vector<int>>& mat, int K) {
int rows=mat.size(),cols=mat[0].size();
vector<vector<int>>dp(rows+1,vector<int>(cols+1,0));
vector<vector<int>>res(rows,vector<int>(cols,0));
for(int i=1;i<=rows;i++){
for(int j=1;j<=cols;j++){
dp[i][j]=dp[i-1][j]+dp[i][j-1]-dp[i-1][j-1]+mat[i-1][j-1];
}
}
for(int i=0;i<rows;i++){
for(int j=0;j<cols;j++){
int x1=max(i-K,0);
int y1=max(j-K,0);
int x2=min(rows-1,i+K);
int y2=min(cols-1,j+K);
res[i][j]=dp[x2+1][y2+1]-dp[x2+1][y1]-dp[x1][y2+1]+dp[x1][y1];
}
}
return res;
}
};
Python:
class Solution:
def matrixBlockSum(self, mat: List[List[int]], K: int) -> List[List[int]]:
m,n=len(mat),len(mat[0])
dp=[[0]*(n+1) for _ in range(m+1)]
res=[[0]*(n) for _ in range(m)]
for i in range(1,m+1):
for j in range(1,n+1):
dp[i][j]=dp[i-1][j]+dp[i][j-1]-dp[i-1][j-1]+mat[i-1][j-1]
for i in range(m):
for j in range(n):
x1=max(i-K,0)
y1=max(j-K,0)
x2=min(m-1,i+K)
y2=min(n-1,j+K)
res[i][j]=dp[x2+1][y2+1]-dp[x2+1][y1]-dp[x1][y2+1]+dp[x1][y1]
return res