从numpy矩阵中查找值列表的按行索引

Question

有没有一种聪明的方法可以使用numpy在矩阵中查找值列表的索引？我们总是可以用迭代的方式来做这件事，但是有没有什么快速可用的方法呢？

我们有一个矩阵：

[[1,2,3],[4,5,6],[7,8,9],[10,11,12]]

和一个列表：

[3,4,7,11]

这里我们需要的是矩阵对应行中的3,4,7和11的索引，即[2,0,0,1]。

这显然是一个用迭代方法编写的简单代码，但我们正在寻找任何现成的实现。

提前谢谢。

Answer 1

对于您的示例，这是可行的：

In [17]: import numpy as np

In [18]: a = np.array([[1,2,3],[4,5,6],[7,8,9],[10,11,12]])

In [19]: l = [3,4,7,11]

In [20]: np.where(a == np.array(l)[:, None])[1]
Out[20]: array([2, 0, 0, 1])

Answer 2

您可以使用numpy.argwhere获得所需的索引

import numpy as np

arr = np.array([[1,2,3],[4,5,6],[7,8,9],[10,11,12]])
search_list = [3,4,7,11]
index_arr = [np.argwhere(arr==x).flatten()[1] for x in search_list]
#index_arr [2, 0, 0, 1]

Answer 3

这是@Akavall解决方案的复制安全版本。如果某个值多次出现，则返回第一个索引：

a = [[1,2,3],[4,5,6],[7,8,9],[10,11,12]]
b = [3,4,7,11]

# introduce a repeated value 
a[1][1] = a[1][0]
a
# [[1, 2, 3], [4, 4, 6], [7, 8, 9], [10, 11, 12]]

i, j = np.where(a == np.c_[b])
j[i.searchsorted(range(len(b)))]
# array([2, 0, 0, 1])