NumPy - 返回数组中元素的第一个索引

Question

什么是在一个ndarray变量和元素之间返回第一个匹配索引的最快方法？我看到很多人使用numpy.where，但那会返回所有索引。

match = 5000
zArray = np.array([[0,1200,200],[1320,24,5000],[5000,234,5230]])

>array([[   0, 1200,  200],
   [1320,   24, 5000],
   [5000,  234, 5230]])

numpy.where(zArray==match)
>(array([1, 2], dtype=int64), array([2, 0], dtype=int64))

我希望返回的是第一个索引，即只有[1,2]。但是numpy.where返回了[1,2]和[2,0]两个结果。

- kabammi

---

你会如何定义“第一个”匹配？按行优先还是按列优先？ - ZdaR

{btsdaf} - kabammi

{btsdaf} - user707650

{btsdaf} - Divakar

Answer 1

您可以使用np.argwhere来获取匹配的索引，以2D数组的形式打包，每行包含每个匹配的索引，然后索引到第一行，如下所示-

np.argwhere(zArray==match)[0]

另外，还有一种更快的方法，使用argmax来获取展平版本上第一个匹配项的索引，以及np.unravel_index用于每个维度的索引元组 -

np.unravel_index((zArray==match).argmax(), zArray.shape)

运行示例 -

In [100]: zArray
Out[100]: 
array([[   0, 1200, 5000], # different from sample for a generic one
       [1320,   24, 5000],
       [5000,  234, 5230]])

In [101]: match
Out[101]: 5000

In [102]: np.argwhere(zArray==match)[0]
Out[102]: array([0, 2])

In [103]: np.unravel_index((zArray==match).argmax(), zArray.shape)
Out[103]: (0, 2)

运行时测试 -

In [104]: a = np.random.randint(0,100,(1000,1000))

In [105]: %timeit np.argwhere(a==50)[0]
100 loops, best of 3: 2.41 ms per loop

In [106]: %timeit np.unravel_index((a==50).argmax(), a.shape)
1000 loops, best of 3: 493 µs per loop

- Divakar

---

{btsdaf} - hpaulj

谢谢Divakar。这很有效，非常适合我的需求。 - kabammi