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

9ceoxa92  于 2023-10-19  发布在  其他
关注(0)|答案(3)|浏览(132)

我想获取一个二维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]))

我需要index array([3,15])。

ffvjumwh

ffvjumwh1#

你需要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)
rnmwe5a2

rnmwe5a22#

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
axzmvihb

axzmvihb3#

使用我创建的numpy_indexed包,您可以简单地编写:

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

相关问题