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

Question

我的矩阵很简单，例如:

# python3 numpy
>>> A
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
array([[0., 0., 0., 0.]])

我需要在 A 中找到一个与 P (1x4) 大小相同的全零区域(一个就足够了)。 所以正确的答案包括:

(2, 0)  # The vertex coordinates of the all-zero rectangular region that P can be matched
(2, 1)
(3, 0)
(3, 1)
(4, 0)
(4, 1)
# Just get any 1 answer

实际上我的A矩阵会达到30,000*30,000的大小。我担心如果写成循环语句会很慢。有什么快速的方法吗？

P的大小不确定，从10*30到4000*80。同时A矩阵缺乏规律性，从任意点开始循环可能需要遍历整个矩阵才能成功匹配

Answer 1

正如@Julien在评论中指出的，一般来说，我们可以使用滑动窗口来完成此类任务。

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
                                                       ^^^^^^^^

作为替代方案，我认为使用 Summed-area table是个好主意。

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

由于我们不需要计算任何给定位置的总和，因此我调整了这个想法并将其实现为仅需要单行缓存。

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

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