【算法思考记录】【前缀和，Python3】力扣1277. 统计全为 1 的正方形子矩阵

原题链接

---

基于前缀和算法的矩阵正方形子矩阵统计解决方案

问题描述

1277号问题要求我们在一个由0和1组成的矩阵中，计算所有完全由1组成的正方形子矩阵的数量。例如，在以下矩阵中：

[
  [0,1,1,1],
  [1,1,1,1],
  [0,1,1,1]
]

我们需要统计所有由1组成的正方形的总数。

解决方案

解决此问题的关键在于有效地计算特定区域内1的数量。为此，我们使用一种称为前缀和的技术。前缀和是一种数据预处理方法，能够快速求解数组中某一段区间的和。

前缀和算法的基本原理

前缀和算法是一种用于快速求解数组或矩阵中某个区间内元素和的技术。这种方法特别适用于需要多次对同一数据集的不同区间求和的情况。它通过预处理数据来优化求和的效率。

一维前缀和

假设我们有一个一维数组 arr，其前缀和数组 preSum 可以这样计算：

1. 初始化 preSum[0] = 0。
2. 对于 arr 中的每个元素 arr[i]，计算 preSum[i+1] = preSum[i] + arr[i]。

这样，preSum[i] 存储了 arr 中从第一个元素到第 i-1 个元素的总和。如果我们想要计算 arr 中从第 l 个元素到第 r 个元素的和，我们只需要计算 preSum[r+1] - preSum[l]。

二维前缀和

二维前缀和是一维前缀和的扩展，用于处理矩阵。假设我们有一个二维数组 matrix，其对应的二维前缀和矩阵 preSum 可以这样计算：

1. 初始化 preSum 的所有元素为0。
2. 对于 matrix 中的每个元素 matrix[i][j]，计算 preSum[i+1][j+1] 的值。这个值等于其左侧元素的前缀和、上方元素的前缀和和左上方元素的前缀和的总和，再减去左上方元素的前缀和，最后加上 matrix[i][j] 本身的值。

计算矩阵中某个子矩形区域内元素的总和时，设该区域的左上角坐标为 (r1, c1)，右下角坐标为 (r2, c2)，则该区域内元素的总和可以表示为：

preSum[r2][c2] - preSum[r1-1][c2] - preSum[r2][c1-1] + preSum[r1-1][c1-1]

这个公式的含义是：从整个矩阵的起点到 (r2, c2) 的总和，减去上方和左方的部分重叠的区域，再加上左上角重复减去的部分。

应用

在解决我们的问题时，我们利用了二维前缀和算法来快速计算矩阵中任意正方形区域内的元素总和。通过这种方式，我们避免了对每个可能的正方形区域进行繁琐的重复计算，显著提高了算法的效率。

代码实现

以下是实现这一算法的Python代码：

class Solution:
    def countSquares(self, matrix: List[List[int]]) -> int:
        row_len, col_len = len(matrix), len(matrix[0])
        pre_sum = [[0] * (col_len + 1) for _ in range(row_len + 1)]

        # 构建前缀和矩阵
        for i, row in enumerate(matrix):
            for j, ele in enumerate(row):
                pre_sum[i + 1][j + 1] = pre_sum[i][j + 1] + pre_sum[i + 1][j] - pre_sum[i][j] + ele

        side = min(row_len, col_len)
        ans = 0

        # 遍历所有可能的正方形
        for k in range(1, side + 1):
            for i in range(k, row_len + 1):
                for j in range(k, col_len + 1):
                    r1, c1 = i - k, j - k
                    r2, c2 = i, j
                    all_one = k * k
                    actual_one = pre_sum[r2][c2] - pre_sum[r1][c2] - pre_sum[r2][c1] + pre_sum[r1][c1]
                    if actual_one == all_one:
                        ans += 1
        return ans

解析

1. 构建前缀和矩阵：首先，我们初始化一个比原矩阵多一行一列的前缀和矩阵，然后逐个计算其值。每个元素的前缀和等于其左侧、上方和左上方元素的前缀和的总和减去左上角的前缀和，再加上当前元素的值。

2. 遍历所有可能的正方形：考虑所有可能的正方形边长，从1到矩阵的最小维度。对于每种边长，我们检查矩阵中的每个位置，看看是否可以形成一个全为1的正方形。

3. 计算并比较实际值和期望值：对于每个可能的正方形位置，我们利用前缀和矩阵快速计算出该正方形内1的实际数量，并将其与该正方形理论上应该有的1的数量（即边长的平方）进行比较。如果两者相等，则找到一个有效的正方形，累加至最终答案中。

这种方法的优势在于减少了重复的计算，使得算法的时间复杂度显著降低，从而能够高效地处理大型矩阵。