题目简介

目标：找出给定二进制矩阵中所有只由 1 组成且面积为 K 的矩阵的数量。

输入：一个 m × n 的二进制矩阵 mat，以及整数 K。

输出：矩阵 mat 中只由 1 组成且面积为 K 的矩阵的数量。

解题思路

本题可以考虑利用滑动窗口算法来解决。具体来说，从左上角开始，每次向右移动一格，如到达矩阵末尾，则将当前窗口向下移动一格，并将当前窗口重新定位到当前行的最左位置。

在每个位置上，需要计算当前窗口内只由 1 组成的矩阵的数量，并求和。

为了加速计算，可以使用前缀和技巧，将矩阵中每一行前缀和缓存下来，这样计算当前窗口内的求和将变得非常快。

代码实现

下面是代码实现，其中 mat_prefix 中存储了矩阵的每一行前缀和。

def count_k_rectangles(mat: List[List[int]], K: int) -> int:
    rows, cols = len(mat), len(mat[0])
    mat_prefix = [[0] * cols for _ in range(rows)]
    for i in range(rows):
        for j in range(cols):
            mat_prefix[i][j] = mat[i][j] + (mat_prefix[i][j - 1] if j > 0 else 0)
    
    cnt = 0
    for i in range(rows):
        for j in range(cols):
            for h in range(i, rows):
                for w in range(j, cols):
                    area = (h - i + 1) * (w - j + 1)
                    if area == K:
                        ones = sum([mat_prefix[k][w] - (mat_prefix[k][j - 1] if j > 0 else 0) for k in range(i, h + 1)])
                        if ones == K:
                            cnt += 1
                    elif area > K:
                        break
    return cnt

示例

mat = [
    [1, 1, 1],
    [1, 1, 0],
    [1, 0, 1]
]
K = 4
print(count_k_rectangles(mat, K))

输出结果为 3，这是因为矩阵中存在三个只由 1 组成、面积为 4 的矩阵，分别是：

1 1
1 1

1 1
   1
1  1

1 1 1
1 1 1