python - 基于单个数组逐行选择元素

Question

假设我有一个大小为 (N,T) 的数组 d，我需要使用形状的 index 从其中选择元素(N,)，其中第一个元素对应于第一行中的索引，等等...我该怎么做？

例如

>>> d
Out[748]: 
array([[ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9],
       [ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9],
       [ 1,  2,  3,  4,  5,  6,  7,  8,  9, 10]])
>>> index
Out[752]: array([5, 6, 1], dtype=int64)

预期输出:

array([[5],
       [6],
       [2])

这是一个数组，包含第一行的第五个元素、第二行的第六个元素和第三行的第二个元素。

更新

由于我将拥有足够大的 N，因此我对更高 N 的不同方法的速度感兴趣。使用N = 30000:

>>> %timeit np.diag(e.take(index2, axis=1)).reshape(N*3, 1)
1 loops, best of 3: 3.9 s per loop
>>> %timeit e.ravel()[np.arange(e.shape[0])*e.shape[1]+index2].reshape(N*3, 1)
1000 loops, best of 3: 287 µs per loop

最后，您建议reshape()。因为我想让它尽可能通用(不知道 N)，所以我使用 [:,np.newaxis] - 它似乎增加了 的持续时间287μs 到 288μs，我会采用:)

Answer 1

这可能很难看，但更有效:

>>> d.ravel()[np.arange(d.shape[0])*d.shape[1]+index]
array([5, 6, 2])

编辑

正如 @deinonychusaur 所指出的，上面的语句可以写得很干净:

d[np.arange(index.size),index]