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

Question

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

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

我想返回一个数组元组，每个数组都包含 yie 中 x 的每个元素的索引

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

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

标签： pythonnumpy

Answer 1

一种解决方案是map

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)