python - numpy 数组中多个元素的索引

Question

我有一个 numpy 数组和一个列表，如下所示

y=np.array([[1],[2],[1],[3],[1],[3],[2],[2]])
x=[1,2,3]

我想返回一个数组元组，每个数组包含 y 中 x 的每个元素的索引。即

(array([[0,2,4]]),array([[1,6,7]]),array([[3,5]]))

这是否可以以矢量化方式完成(没有任何循环)？

Answer 1

一种解决方案是映射

y = y.reshape(1,len(y))
map(lambda k: np.where(y==k)[-1], x)

[array([0, 2, 4]), 
 array([1, 6, 7]), 
 array([3, 5])]

<小时/>

合理的性能。对于 100000 行，

%timeit list(map(lambda k: np.where(y==k), x))
3.1 ms ± 113 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)