在二维numpy数组中查找匹配行

Question

我想要得到一个匹配行的二维Numpy数组的索引。例如，我的数组是这样的：

vals = np.array([[0, 0],
                 [1, 0],
                 [2, 0],
                 [0, 1],
                 [1, 1],
                 [2, 1],
                 [0, 2],
                 [1, 2],
                 [2, 2],
                 [0, 3],
                 [1, 3],
                 [2, 3],
                 [0, 0],
                 [1, 0],
                 [2, 0],
                 [0, 1],
                 [1, 1],
                 [2, 1],
                 [0, 2],
                 [1, 2],
                 [2, 2],
                 [0, 3],
                 [1, 3],
                 [2, 3]])

我想要得到匹配行[0, 1]的索引，它们是索引3和15。 当我执行类似于numpy.where(vals == [0 ,1])的操作时，我会得到...

(array([ 0,  3,  3,  4,  5,  6,  9, 12, 15, 15, 16, 17, 18, 21]), array([0, 0, 1, 1, 1, 0, 0, 0, 0, 1, 1, 1, 0, 0]))

我想要索引数组([3, 15])。

- b10hazard

Answer 1

使用np.where函数可获取索引：

>>> np.where((vals == (0, 1)).all(axis=1))
(array([ 3, 15]),)

或者，正如文档所述：

如果只给出条件，则返回condition.nonzero()

您可以直接对.all返回的数组调用.nonzero()：

>>> (vals == (0, 1)).all(axis=1).nonzero()
(array([ 3, 15]),)

为了拆解它：

>>> vals == (0, 1)
array([[ True, False],
       [False, False],
       ...
       [ True, False],
       [False, False],
       [False, False]], dtype=bool)

并在该数组上调用.all方法（使用axis=1），该方法将返回True，其中两个布尔值为True：

>>> (vals == (0, 1)).all(axis=1)
array([False, False, False,  True, False, False, False, False, False,
       False, False, False, False, False, False,  True, False, False,
       False, False, False, False, False, False], dtype=bool)

获取哪些索引是 True：

>>> np.where((vals == (0, 1)).all(axis=1))
(array([ 3, 15]),)

或者

>>> (vals == (0, 1)).all(axis=1).nonzero()
(array([ 3, 15]),)

我认为我的解决方案更易读一些，但正如unutbu所指出的那样，以下方法可能更快，并且与 (vals == (0, 1)).all(axis=1) 返回相同的值：

>>> (vals[:, 0] == 0) & (vals[:, 1] == 1)

- Russia Must Remove Putin

---

我倾向于使用np.nonzero而不是np.where的别名，以避免与完全不同的np.where(bool, if_true, if_false)函数混淆。 - Eric

Answer 2

In [5]: np.where((vals[:,0] == 0) & (vals[:,1]==1))[0]
Out[5]: array([ 3, 15])

我不确定为什么，但这种方法比 np.where((vals == (0, 1)).all(axis=1)) 快得多：

In [34]: vals2 = np.tile(vals, (1000,1))

In [35]: %timeit np.where((vals2 == (0, 1)).all(axis=1))[0]
1000 loops, best of 3: 808 µs per loop

In [36]: %timeit np.where((vals2[:,0] == 0) & (vals2[:,1]==1))[0]
10000 loops, best of 3: 152 µs per loop

- unutbu

Answer 3

使用我创建的numpy_indexed软件包，您只需简单地编写：

import numpy_indexed as npi
print(np.flatnonzero(npi.contains([[0, 1]], vals)))

- Eelco Hoogendoorn

---

很高兴你做了这个。你应该明确表明你的隶属关系以符合网站规则。我注意到你在最近我的研究中回答了另一个问题并做到了这一点。 - Mad Physicist

你能展示一些基准测试吗？否则我可能无法说服自己或其他人使用它。 - Jiadong