如何快速在大型numpy矩阵中找到首个匹配的全零子矩阵？

问题场景 #

现有Numpy矩阵示例：

import numpy as np

A = np.array([[0., 0., 1., 1., 1.],
              [0., 0., 1., 1., 1.],
              [0., 0., 0., 0., 0.],
              [0., 0., 0., 0., 0.],
              [0., 0., 0., 0., 0.]])
P = np.array([[0., 0., 0., 0.]])

需要在A中找到一个与P尺寸（h_p × w_p，示例中为1×4）完全相同的全零矩形区域，返回左上角顶点坐标即可（示例中合法坐标如(2,0)、(3,1)等）。实际场景中A可达30000×30000，P尺寸范围为10×30到4000×80，需避免低效的全矩阵遍历。

高效实现方案 #

方法1：二维前缀和快速检测 #

利用前缀和可以O(1)计算任意矩形区域的元素和，若和为0则说明该区域全为0，找到第一个匹配后立即终止遍历。

代码实现：

def find_zero_rect(A, P):
    h_p, w_p = P.shape
    h_a, w_a = A.shape
    if h_p > h_a or w_p > w_a:
        return None  # P尺寸超过A，无匹配可能
    
    # 计算二维前缀和（左开右闭索引，方便区域和计算）
    prefix = np.zeros((h_a + 1, w_a + 1), dtype=np.float64)
    prefix[1:, 1:] = np.cumsum(np.cumsum(A, axis=0), axis=1)
    
    # 遍历所有可能的左上角坐标，找到第一个符合条件的就返回
    for x in range(h_a - h_p + 1):
        for y in range(w_a - w_p + 1):
            # 计算当前矩形区域的元素和
            sum_val = prefix[x+h_p, y+w_p] - prefix[x, y+w_p] - prefix[x+h_p, y] + prefix[x, y]
            if np.isclose(sum_val, 0):  # 兼容浮点数精度问题
                return (x, y)
    return None

方法2：行内零段预筛选+跨行匹配（更高效） #

先逐行筛选出长度≥P列数的连续零段，再通过滑动窗口检查连续P行中是否存在重叠的有效零段，大幅减少需要检查的坐标数量。

代码实现：

def find_zero_rect_fast(A, P):
    h_p, w_p = P.shape
    h_a, w_a = A.shape
    if h_p > h_a or w_p > w_a:
        return None
    
    # 预处理每一行：提取所有长度≥w_p的连续零段，存储为(start, end)（end为区间右边界，左闭右开）
    row_segments = []
    for row in A:
        zeros = np.where(row == 0)[0]
        if len(zeros) == 0:
            row_segments.append([])
            continue
        # 分割连续零的片段
        split_points = np.where(np.diff(zeros) > 1)[0] + 1
        segments = np.split(zeros, split_points)
        # 筛选有效片段
        valid_segs = []
        for seg in segments:
            if len(seg) >= w_p:
                valid_segs.append((seg[0], seg[-1] + 1))
        row_segments.append(valid_segs)
    
    # 滑动窗口检查连续h_p行的共同有效区间
    for x_start in range(h_a - h_p + 1):
        window_segs = row_segments[x_start:x_start+h_p]
        # 窗口内任意一行无有效段则跳过
        if any(len(segs) == 0 for segs in window_segs):
            continue
        # 逐步求多行间的重叠有效区间
        current_overlap = window_segs[0][0]
        for segs in window_segs[1:]:
            new_overlap = None
            for s in segs:
                overlap_start = max(current_overlap[0], s[0])
                overlap_end = min(current_overlap[1], s[1])
                if overlap_end - overlap_start >= w_p:
                    new_overlap = (overlap_start, overlap_end)
                    break
            if new_overlap is None:
                break
            current_overlap = new_overlap
        if current_overlap is not None:
            return (x_start, current_overlap[0])
    return None

性能说明 #

• 方法1的前缀和计算为numpy向量化操作，遍历部分在找到第一个匹配后立即终止，实际性能优于纯循环，时间复杂度O(h_a*w_a)。
• 方法2通过预筛选减少了无效遍历，当A中零区域占比较小时，性能优势明显；numpy的np.where等向量化操作可高效处理大矩阵的行内零段检测。
• 若A为浮点数矩阵，需用np.isclose判断区域和是否为0，避免精度误差。

内容的提问来源于stack exchange，提问作者ojipadeson