python - 确定数组 B 的每个元素在另一个数组 A 中的位置

Question

给定两个列表:

In [518]: A
Out[518]: [3, 4, 2, 1, 7, 6, 5]

In [519]: B
Out[519]: [4, 6]

B 中的每个元素都存在于 A 中，无一异常(exception)。

我想检索 A 中所示的 B 的索引数组。例如，4 出现在 A 中的索引 1 中，而 6 出现在 B 的位置 5 中。对于这种情况，我的预期输出是 [1, 5]。

这就是我为获取索引所做的事情:

In [520]: np.flatnonzero(np.in1d(a, b))
Out[520]: array([1, 5])

不幸的是，这在大多数其他情况下不起作用。例如，如果 B = [6, 4]，我的方法在应该输出 [5, 1] 时仍然输出 [1, 5] >。

有没有一种有效的 numpy 方法来实现我想要实现的目标？

Answer 1

IIUC:

In [71]: a
Out[71]: array([3, 4, 2, 1, 7, 6, 5, 6, 4])

In [72]: b
Out[72]: array([4, 6])

In [73]: np.where(a==b[:,None])[1]
Out[73]: array([1, 8, 5, 7], dtype=int64)

In [74]: b = np.array([6, 4])

In [75]: np.where(a==b[:,None])[1]
Out[75]: array([5, 7, 1, 8], dtype=int64)

更新:如果您只需要第一次出现的索引(以防A数组中存在重复项)，则使用this solution from @Divakar ，这会更快:

In [84]: (a==b[:,None]).argmax(1)
Out[84]: array([5, 1], dtype=int64)