python - 在多维numpy数组中查找多个查询的索引

Question

我正在寻找一种在多维数组中查找查询数组索引的方法。例如:

arr = np.array([[17,  5, 19,  9],
   [18, 13,  3,  7],
   [ 8,  1,  4,  2],
   [ 8,  9,  7, 19],
   [ 6, 11,  8,  5],
   [11, 16, 13, 18],
   [ 0,  1,  2,  9],
   [ 1,  7,  4,  6]])

我可以找到一个查询的索引:

np.where(arr==1)
# (array([2, 6, 7]), array([1, 1, 0]))

是否有任何 numpy 解决方案可以对多个值执行此操作，以替换以下 for 循环？

for q in queries:
    np.where(arr==q)

如果数组和查询都是一维的，我可以使用 searchsorted 作为 this answer但它不适用于多维数组。

Answer 1

IIUC你可以尝试这个:

In[19]:np.where((arr==4)|(arr==5))
Out[19]: (array([0, 2, 4, 7], dtype=int64), array([1, 2, 3, 2], dtype=int64))