从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]。

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

提前致谢。

- SamGhatak

---

你可以使用Numpy Argwhere在列表推导式中获取numpy数组中元素的位置（如果你想搜索多个元素）。 - Sohaib Farooqi

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])

- Akavall

---

不会帮助查找属于同一子数组的元素，因为它只能找到第一个元素。 - newbie

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]

- Sohaib Farooqi

---

运行良好，但我尝试不使用for循环，而是让框架在可能的情况下执行任何优化。 - SamGhatak

@SamGhatak 有道理 :) - Sohaib Farooqi

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])

- Paul Panzer