Python中加速子矩阵索引搜索的高效方法寻求

原有的双重循环逐块比对方法，在矩阵规模较大时效率极低——它需要逐个提取子矩阵并逐元素校验，时间复杂度为O((H-h+1)*(W-w+1)hw)（H、W为主矩阵高宽，h、w为子矩阵高宽）。下面两种基于numpy/scipy的优化方案，能大幅提升搜索速度：

方案一：利用numpy滑动窗口批量比对 #

通过numpy.lib.stride_tricks.as_strided一次性生成主矩阵的所有候选子窗口，再借助numpy的广播机制完成批量元素比对，把Python层面的循环转移到底层C实现，速度提升明显。

import numpy as np

def fast_search_submatrix(kernel, array):
    arr_h, arr_w = array.shape
    ker_h, ker_w = kernel.shape
    
    # 定义滑动窗口的形状与步长
    window_shape = (ker_h, ker_w)
    strides = array.strides + array.strides
    
    # 生成所有符合尺寸的子窗口（无内存拷贝）
    windows = np.lib.stride_tricks.as_strided(
        array,
        shape=(arr_h - ker_h + 1, arr_w - ker_w + 1, ker_h, ker_w),
        strides=strides
    )
    
    # 批量校验所有窗口，找到完全匹配的位置
    matches = (windows == kernel).all(axis=(2, 3))
    y_indices, x_indices = np.where(matches)
    
    if len(x_indices) > 0:
        print(f"子矩阵在x={x_indices[0]}, y={y_indices[0]}位置找到。")
        return [x_indices[0], y_indices[0]]
    print("主矩阵中未找到该子矩阵。")

方案二：基于二维互相关的快速匹配 #

将子矩阵匹配转化为二维互相关计算：当互相关结果等于子矩阵元素总和时，大概率是匹配位置（需额外校验避免元素和冲突）。scipy的correlate2d采用优化的卷积算法，时间复杂度更低，适合大矩阵场景。

import numpy as np
from scipy.signal import correlate2d

def conv_based_search(kernel, array):
    ker_h, ker_w = kernel.shape
    # 计算主矩阵与子矩阵的二维互相关
    corr = correlate2d(array, kernel, mode='valid')
    # 子矩阵元素总和，匹配时互相关结果等于该值
    ker_total = kernel.sum()
    
    # 筛选出候选位置并逐一校验
    y_candidates, x_candidates = np.where(corr == ker_total)
    for y, x in zip(y_candidates, x_candidates):
        if (array[y:y+ker_h, x:x+ker_w] == kernel).all():
            print(f"子矩阵在x={x}, y={y}位置找到。")
            return [x, y]
    print("主矩阵中未找到该子矩阵。")

性能说明 #

• 滑动窗口方案：子矩阵较小时，批量比对的优势最明显，几乎没有额外开销；
• 互相关方案：矩阵规模越大，相比循环方法的速度提升越显著，适合处理大型图像或数据集。

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