在二维 numpy 数组中查找匹配的行
Posted
技术标签:
【中文标题】在二维 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)))
【讨论】:
很高兴你做了这个。您应该明确您的隶属关系以遵守网站规则。我注意到您最近在我的研究中提出的另一个答案中做了。以上是关于在二维 numpy 数组中查找匹配的行的主要内容,如果未能解决你的问题,请参考以下文章