python - 查找较大二维矩阵中存在较小二维矩阵的位置

Question

假设我有以下 6x6 矩阵 a 和 2x2 矩阵 b:

a = np.array([[1, 1, 1, 0, 0, 0],
              [1, 0, 1, 0, 1, 0],
              [1, 0, 0, 0, 0, 1],
              [1, 1, 1, 1, 1, 0],
              [0, 0, 0, 0, 0, 0],
              [1, 0, 0, 1, 1, 1]])

b = np.array([[0, 1],
              [0, 0]])

我想创建一个真值表，其中较小的数组二维矩阵存在于较大的数组内部，但仅选择左上角索引

预期输出:

>>> array([[False, False, False, False, False, False],
           [False, True,  False, True,  False, False],
           [False, False, False, False, False, False],
           [False, False, False, False, False, False],
           [False, False, False, False, False, False],
           [False, False, False, False, False, False]])

我想我可以构建一个巨大的循环并使用 np.roll() 进行评估，但我希望有一个优雅的解决方案？

Answer 1

您可以使用np.lib.stride_tricks.as_strided创建带有两个附加维度的 a 版本用于比较。然后您可以使用 .all(axis=(-2, -1)) 来验证相等性:

In [3]: a2 = np.lib.stride_tricks.as_strided(a, shape=a.shape + b.shape, strides=2*a.strides)
In [4]: a2 = a2[:-b.shape[0]+1, :-b.shape[1]+1]
In [5]: (a2 == b).all(axis=(-2, -1))                                                          
Out[5]: 
array([[False, False, False, False, False],
       [False,  True, False,  True, False],
       [False, False, False, False, False],
       [False, False, False, False, False],
       [False, False, False, False, False]])

如果你想要额外的填充，你可以使用 np.pad事先:

a = np.pad(a, ((0, 1), (0, 1)), mode='wrap')

然后按照上面的 as_strided 继续操作。