如何在numpy中使用列值获取某行的索引？

Question

如果我有下面的数组arr.Arr.A，那么我就可以用这个数组来获取第一行的索引。

>>> arr
array([[ 0,  1,  2,  3,  4],
       [ 5,  6,  7,  8,  9],
       [10, 11, 12, 13, 14],
       [ 0,  1,  2,  3,  4],
       [ 5,  6,  7,  8,  9],
       [10, 11, 12, 13, 14]])

现在如果我想获取第一行，我会这样做。

>>> arr[0]
array([0, 1, 2, 3, 4])

然而，当我使用np.where来定位某一行时，例如。

>>> np.where(arr == [0,1,2,3,4])

我得到的结果是这样的!

(array([0, 0, 0, 0, 0, 3, 3, 3, 3, 3], dtype=int64), 
array([0, 1, 2, 3, 4, 0, 1, 2, 3, 4], dtype=int64))

然而，这并不是我想要的。我想得到行的索引，例如：。

(array([0, 3], dtype=int64)

有什么方法可以实现吗? 任何建议都非常感激!

Answer 1

我想你是想检查这些行是否等于一个给定数组。在这种情况下，你需要 all:

np.where((arr == [0,1,2,3,4]).all(1))
# (array([0, 3]),)