python - 根据另一个数组的索引列表从一个数组中获取元素

Question

我有一个像这样的 2 numpy 数组

a = [array([ 0.1,  0.1,  0.1]), array([ 0.2,  0.2,  0.2])]

b = [0 0 0 1]

我想要的是这样的--

c = [[0.1,  0.1,  0.1],[0.1,  0.1,  0.1],[0.1,  0.1,  0.1],[0.2, 0.2, 0.2]]

即a 的元素基于 b 的索引。

有没有一种方法可以使用 numpy 和矢量化来实现这一点，即无需循环值？

Answer 1

如果将 a 存储为二维 numpy 数组:

>>> a = np.array([[0.1, 0.1, 0.1], [0.2, 0.2, 0.2]])
# result: array([[ 0.1,  0.1,  0.1],
#                [ 0.2,  0.2,  0.2]])

甚至通过 a = np.array(a) 将 a 转换为 numpy 数组，然后您可以使用列表 b 根据需要访问元素:

>>> b = [0,0,0,1]
>>> print(a[b])
array([[ 0.1,  0.1,  0.1],
       [ 0.1,  0.1,  0.1],
       [ 0.1,  0.1,  0.1],
       [ 0.2,  0.2,  0.2]])

如果你需要一个 list 作为输出，那么使用 numpy 数组的 tolist() 方法:

>>> (np.asarray(a)[b]).tolist()
[[0.1, 0.1, 0.1], [0.1, 0.1, 0.1], [0.1, 0.1, 0.1], [0.2, 0.2, 0.2]]