这个题用brunt force的方法时间复杂度是O(mnK^2),题目的测试case 里K是可以特别大的。想去掉K的话,做法是定义一个2d array cumsum, 使得cumsum[i][j] = A[:i][:j].
用dp的方法来获得这个cumsum,之后就好做了
class Solution:
def matrixBlockSum(self, mat: List[List[int]], K: int) -> List[List[int]]:
#TC: O(m*n)
#SC: O(m*n) m = len(mat); n = len(mat[0])
m = len(mat)
n = len(mat[0])
cumsum = [[0]*n for _ in range(m)]
cumsum[0][0] = mat[0][0]
for row in range(1, m):
cumsum[row][0] = cumsum[row-1][0] + mat[row][0]
for col in range(1, n):
cumsum_col = 0
for row in range(m):
cumsum[row][col] = cumsum[row][col-1] + cumsum_col + mat[row][col]
cumsum_col += mat[row][col]
res = [[0]*n for _ in range(m)]
for row in range(m):
for col in range(n):
start_row = row - K - 1
start_col = col - K - 1
end_row = min(row + K, m-1)
end_col = min(col + K, n-1)
if start_row < 0:
a = 0
else:
a = cumsum[start_row][end_col]
if start_col < 0:
b = 0
else:
b = cumsum[end_row][start_col]
if start_row < 0 or start_col < 0:
c = 0
else:
c = cumsum[start_row][start_col]
res[row][col] = cumsum[end_row][end_col] - a - b + c
return res