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

Posted 2023-03-12

【中文标题】在二维 numpy 数组中查找匹配的行

【英文标题】：Find matching rows in 2 dimensional numpy array

【发布时间】：2014-11-07 12:53:29

【问题描述】：

我想获取与行匹配的二维 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])。

【参考方案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)

【讨论】：

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

【参考方案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

【讨论】：

这个方法可以让我分别匹配两个轴，这正是我所需要的。

【参考方案3】：

使用numpy_indexed包，你可以简单地写：

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

【讨论】：

很高兴你做了这个。您应该明确您的隶属关系以遵守网站规则。我注意到您最近在我的研究中提出的另一个答案中做了。