leetcode1314

 1 class Solution:
 2     def matrixBlockSum(self, mat: ‘List[List[int]]‘, K: int) -> ‘List[List[int]]‘:
 3         r,c = len(mat),len(mat[0])
 4         prefixsum = [[0 for _ in range(c)]for _ in range(r)]
 5         for i in range(r):
 6             prefixsum[i][0] = mat[i][0]
 7             for j in range(1,c):
 8                 prefixsum[i][j] = mat[i][j] + prefixsum[i][j-1]
 9         for j in range(c):
10             for i in range(1,r):
11                 prefixsum[i][j] = prefixsum[i-1][j] + prefixsum[i][j]
12         #print(prefixsum)
13         res = [[0 for _ in range(c)]for _ in range(r)]
14         for i in range(r):
15             for j in range(c):
16                 r_low,r_high = max(i - K,0),min(i + K,r-1)
17                 c_low,c_high = max(j - K,0),min(j + K,c-1)
18                 #print(r_low,c_low,r_high,c_high)
19                 
20                 if r_low >= 1 and c_low >= 1:
21                     res[i][j] = prefixsum[r_high][c_high] - prefixsum[r_low-1][c_high] - prefixsum[r_high][c_low-1] + prefixsum[r_low-1][c_low-1]
22                 elif r_low >= 1 and c_low == 0:
23                     res[i][j] = prefixsum[r_high][c_high] - prefixsum[r_low-1][c_high]
24                 elif r_low == 0 and c_low >= 1:
25                     res[i][j] = prefixsum[r_high][c_high] - prefixsum[r_high][c_low-1]
26                 else:
27                     res[i][j] = prefixsum[r_high][c_high]
28         return res