Python 从矩阵中找到第一个匹配子矩阵的快速方法

正如@Julien在 comments 中指出的那样，一般来说，我们可以使用滑动窗口来完成这类任务.

def find_all_zero_region_by_sliding_window(a, shape):
    x, y = np.nonzero(np.lib.stride_tricks.sliding_window_view(a, shape).max(axis=-1).max(axis=-1) == 0)
    return np.stack((x, y), axis=-1)


find_all_zero_region_by_sliding_window(A, P.shape)

然而，不幸的是，这需要大量内存.

numpy.core._exceptions.MemoryError: Unable to allocate 11.3 TiB for an array with shape (26001, 29921, 4000) and data type float32
                                                       ^^^^^^^^

作为另一种 Select ，我认为使用Summed-area table是一个好主意.

它类似于上面的滑动窗口方法，但我们可以计算和(非常高效)并搜索它为零的位置，而不是寻找最大值. 请注意，这假设A不包含任何负值.否则，您将不得不使用numpy.abs.

因为我们不需要能够计算任何给定位置的总和，所以我采用了这个 idea 并将其实现为只需要单线缓存.

import numpy as np
from typing import Tuple


def find_all_zero_region(arr: np.ndarray, kernel_size: Tuple[int, int]) -> np.ndarray:
    input_height, input_width = arr.shape
    kernel_height, kernel_width = kernel_size

    matches = []

    # Calculate summed_line for y==0.
    summed_line = arr[:kernel_height].sum(axis=0)

    for y in range(input_height - kernel_height + 1):
        # Update summed_line for row y.
        if y != 0:  # Except y==0, which already calculated above.
            # Adding new row and subtracting old row.
            summed_line += arr[y + kernel_height - 1] - arr[y - 1]

        # Calculate kernel_sum for (y, 0).
        kernel_sum = summed_line[:kernel_width].sum()
        if kernel_sum == 0:
            matches.append((y, 0))

        # Calculate kernel_sum for (y, 1) to (y, right-edge).
        # Using the idea of a summed-area table, but in 1D (horizontally).
        (all_zero_region_cols,) = np.nonzero(kernel_sum + np.cumsum(summed_line[kernel_width:] - summed_line[:-kernel_width]) == 0)
        for col in all_zero_region_cols:
            matches.append((y, col + 1))

    if not matches:
        # For Numba, output must be a 2d array.
        return np.zeros((0, 2), dtype=np.int64)
    return np.array(matches, dtype=np.int64)

正如您所看到的，这使用了循环，但它应该比您想象的要快得多，因为所需的内存相对较小，并且计算/比较的次数大大减少. 以下是一些计时代码.

import time


rng = np.random.default_rng(0)
A = rng.integers(0, 2, size=(30000, 30000)).astype(np.float32)

P = np.zeros(shape=(4000, 80))

# Create an all-zero region in the bottom right corner which will be searched last.
A[-P.shape[0] :, -P.shape[1] :] = 0

started = time.perf_counter()
result = find_all_zero_region(A, P.shape)
print(f"{time.perf_counter() - started} sec")
print(result)
# 3.541154200000001 sec
# [[26000 29920]]

此外，使用Numba可以更快地实现此功能. 只需按如下方式添加注释:

import numba


@numba.njit("int64[:,:](float32[:,:],UniTuple(int64,2))")
def find_all_zero_region_with_numba(arr: np.ndarray, kernel_size: Tuple[int, int]) -> np.ndarray:
    ...

started = time.perf_counter()
find_all_zero_region_with_numba(A, P.shape)
print(f"{time.perf_counter() - started} sec")
# 1.6005743999999993 sec

请注意，我实现它是为了查找全零区域的所有位置，但您也可以让它在第一个位置返回. 因为它使用循环，所以平均执行时间会更快.