使用索引数组从 numpy 数组中获取值

Question

我正在尝试找到一个简单的 numpy 命令来根据第二个坐标的索引数组获取 3D 矩阵的值。这可以这样做：

import numpy as np
entry = np.array([[1, 2],[3,4],[5, 6]])
entries = np.stack([entry, entry, entry, entry])
indices = np.array([2, 1, 0, 1])

r = np.array([entries[i, indices[i], :] for i in range(len(indices))])

此代码产生正确的结果，但使用列表理解而不是单个命令。

有没有同样的 numpy 操作？我期望 np.take() 能够工作，但它不起作用，因为它生成形状为 (4, 4, 2) 而不是 (4, 2) 的数组

Answer 1

我发现下面的代码可以生成与示例中的 r 相同的数组。

np.take(entry,indices,axis=0) 

看来用take函数不需要堆叠。