python – 在2维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]).

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

或者,正如文档所述：

If only condition is given, return 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方法(轴= 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)