python - 在 3d 数组中查找 2d 数组

Question

有没有一种快速的方法可以找到 2d 数组在 3d 数组中的所有索引？

我有这个 3d numpy 数组:

arr = np.array([
        [[0,1],[0,2],[0,3],[0,4],[0,4],[0,5],[0,5],[0,5],[0,5],[0,5]],
        [[0,1],[0,2],[0,2],[0,2],[0,3],[0,4],[0,4],[0,4],[0,5],[0,5]],
        [[0,1],[0,2],[0,3],[0,3],[0,3],[0,4],[0,4],[0,5],[0,5],[0,5]]
       ])

我想找到 [0,4] 出现的所有索引。我试过这个:

whereInd = np.argwhere(arr == np.array([0,4]))

但它不起作用。预期结果是:

[[0 3],[0 4],[1 5],[1 6],[1 7],[2 5],[2 6]]

另一个问题是，这样会不会很快？因为我想将它用于 (10000,100,2) 数组。

Answer 1

使用 argwhere() 是个好主意，但您还需要使用 all() 来获得所需的输出:

>>> np.argwhere((arr == [0, 4]).all(axis=2))
array([[0, 3],
       [0, 4],
       [1, 5],
       [1, 6],
       [1, 7],
       [2, 5],
       [2, 6]])

此处all()用于检查比较后的每一行是否为[True, True](即该行等于[0 , 4]).在 3D 数组中，axis=2 指向行。

这会将维数减少为两个，并且 argwhere() 返回所需的索引数组。

---

关于性能，此方法应该相当快地处理您指定大小的数组:

In [20]: arr = np.random.randint(0, 10, size=(10000, 100, 2))
In [21]: %timeit np.argwhere((arr == [0, 4]).all(axis=2))
10 loops, best of 3: 44.9 ms per loop