此题解采用了二维前缀和的方法来计算任意子矩阵的和。首先,通过构造一个前缀和矩阵 pre_sum,使得每个元素 pre_sum[i][j] 代表原矩阵左上角到 (i-1, j-1) 的所有元素之和。这样,任意子矩阵的和可以通过 pre_sum 矩阵在 O(1) 时间内得到,大大提高了查询效率。具体到每个元素 mat[i][j] 的 k 范围内的矩阵和,只需要通过调整边界,使用前缀和矩阵的差分方法即可快速求出。
时间复杂度: O(mn)
空间复杂度: O(mn)
# 定义解决方案类
class Solution:
def matrixBlockSum(self, mat: List[List[int]], k: int) -> List[List[int]]:
m, n = len(mat), len(mat[0])
nummatrix = NumMatrix(mat)
answer = [[0]*n for _ in range(m)]
# 遍历矩阵中的每个元素
for x in range(m):
for y in range(n):
x1, y1 = max(x-k, 0), max(y-k, 0)
x2, y2 = min(x+k, m-1), min(y+k, n-1)
# 使用前缀和矩阵计算子矩阵和
answer[x][y] = nummatrix.sumRegion(x1, y1, x2, y2)
return answer
class NumMatrix:
def __init__(self, matrix):
m, n = len(matrix), len(matrix[0])
self.pre_sum = [[0]*(n+1) for _ in range(m+1)]
# 构造前缀和矩阵
for i in range(1, m+1):
for j in range(1, n+1):
self.pre_sum[i][j] = self.pre_sum[i-1][j] + self.pre_sum[i][j-1] - self.pre_sum[i-1][j-1] + matrix[i-1][j-1]
def sumRegion(self, x1, y1, x2, y2):
# 使用前缀和矩阵差分计算子矩阵和
return self.pre_sum[x2+1][y2+1] - self.pre_sum[x1][y2+1] - self.pre_sum[x2+1][y1] + self.pre_sum[x1][y1]