如何实现函数判断二维numpy数组中是否存在指定精确匹配子二维数组

核心实现逻辑 #

首先需要完成两层判断：

• 先做尺寸边界校验：如果目标子矩阵B的行列数大于原数组A的行列数，直接判定不存在匹配
• 遍历原数组所有可能的、和B尺寸一致的子块，只要有任意一个子块和B完全匹配，就返回True

纯NumPy高效实现（适配NumPy 1.20+） #

利用NumPy内置的滑动窗口视图功能，避免手动写双重循环，性能更高：

import numpy as np

def has_submatrix(A, B):
    # 兼容列表输入，自动转换为NumPy数组
    A = np.asarray(A)
    B = np.asarray(B)

    h_A, w_A = A.shape
    h_B, w_B = B.shape

    # 边界情况：子矩阵尺寸超出原矩阵，直接返回不存在
    if h_B > h_A or w_B > w_A:
        return False

    # 生成所有位置的滑动窗口，输出形状为（行起始位置数, 列起始位置数, 窗口高度, 窗口宽度）
    windows = np.lib.stride_tricks.sliding_window_view(A, (h_B, w_B))

    # 逐窗口判断是否匹配
    for win in windows.reshape(-1, h_B, w_B):
        if np.array_equal(win, B):
            return True
    return False

示例测试 #

按照题目给出的测试用例调用：

A = [
    [0., 0., 1., 0., 0.],
    [0., 0., 0., 1., 1.],
    [0., 0., 1., 1., 1.]
]
B = [
    [1., 1.],
    [1., 1.]
]

print(has_submatrix(A, B))
# 输出 True

低版本NumPy兼容实现 #

如果使用的NumPy版本低于1.20，没有内置滑动窗口接口，可以用双重遍历的方式实现，逻辑更通用：

def has_submatrix_compat(A, B):
    A = np.asarray(A)
    B = np.asarray(B)
    h_A, w_A = A.shape
    h_B, w_B = B.shape
    if h_B > h_A or w_B > w_A:
        return False
    # 遍历所有可能的窗口起始行、列坐标
    for i in range(h_A - h_B + 1):
        for j in range(w_A - w_B + 1):
            # 截取对应子块做匹配判断
            sub_block = A[i:i+h_B, j:j+w_B]
            if np.array_equal(sub_block, B):
                return True
    return False

注意事项 #

• 针对浮点数匹配场景，如果存在计算带来的精度误差（比如实际值为1.0000000002这类非精确值），可以把对比函数从np.array_equal替换为np.allclose(sub_block, B, atol=1e-8)，避免浮点尾差导致的匹配失败
• 大尺寸原数组场景下，使用滑动窗口视图的实现性能远高于双重循环版本

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