python - 在numpy数组中查找子数组的索引

Question

我有两个 numpy 数组，一个更大，一个更小：

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

是否有一个函数可以用来查找较大数组的索引，其中较小的实例中有一个？

理想结果：

instances[0] = [[2, 0], [2, 1]]
instances[1] = [[1, 1], [1,2]]

非常感谢！

标签： pythonarraysnumpyfind

Answer 1

据我所知，没有快速的 numpy 函数可以做到这一点，但您可以快速循环并检查。

def find_instances(a,b):
    instances = []
    for i in range(a.shape[0] - b.shape[0] + 1):
        for j in range(a.shape[1] - b.shape[1] + 1):
            if np.all(a[i:i+b.shape[0], j:j+b.shape[1]] == b):
                instances.append([i,j])
    return instances

这里每个实例都是 a 的左上角与 b 的左上角匹配的点。不是您要求的输出，但如果您真的需要从那里获取其余索引，则很容易获得它们。希望有帮助！