矩阵区域和

标签: 数组 矩阵 前缀和

难度: Medium

给你一个 m x n 的矩阵 mat 和一个整数 k ,请你返回一个矩阵 answer ,其中每个 answer[i][j] 是所有满足下述条件的元素 mat[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

Submission

运行时间: 41 ms

内存: 17.6 MB

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]

Explain

此题解采用了二维前缀和的方法来计算任意子矩阵的和。首先,通过构造一个前缀和矩阵 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]

Explore

使用比原矩阵行和列都多1的矩阵(通常初始化为0),主要是为了方便计算边界情况,避免在进行前缀和计算时需要频繁检查索引边界。这种方法允许我们在不添加额外条件语句的情况下直接使用公式计算任意子矩阵的和,包括当子矩阵触及原矩阵的边界时。具体来说,当子矩阵的起点为(0,0)时,通过引入额外的0行和0列,可以确保不会访问到负索引,从而简化代码的复杂度。

在计算前缀和时,`self.pre_sum[i-1][j]`和`self.pre_sum[i][j-1]`分别包含了`matrix[i-1][j-1]`单元格左边和上边的元素的和,这导致`matrix[i-1][j-1]`被重复计算了一次。因此,需要减去`self.pre_sum[i-1][j-1]`来消除这种重复计算的影响。这个减去的操作确保了每个元素值只被计算一次,从而准确地得出从矩阵原点到当前元素构成的子矩阵的元素总和。

在计算子矩阵和时,通过设置`x1 = max(x-k, 0)`和`y1 = max(y-k, 0)`来确保不会有负数索引,这样可以避免数组越界。同时,通过设置`x2 = min(x+k, m-1)`和`y2 = min(y+k, n-1)`确保索引不会超过矩阵的实际大小。这些边界值的调整确保了即使子矩阵的请求超出了原始矩阵的边界,代码仍然能够正确地处理并返回正确的矩阵区域和,不会导致运行时错误。