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

Question

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

例如

>>> 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])

包含第一行第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]

Answer 2

可能有更好的方法，但take, diag和reshape的组合会做:

In [137]: np.diag(d.take(index, axis=1)).reshape(3, 1)
Out[137]: 
array([[5],
       [6],
       [2]])

编辑

与@Emanuele Paolinis的替代比较，添加重塑以匹配搜索的输出:

In [142]: %timeit d.reshape(d.size)[np.arange(d.shape[0])*d.shape[1]+index].reshape(3, 1)
100000 loops, best of 3: 9.51 µs per loop
In [143]: %timeit np.diag(d.take(index, axis=1)).reshape(3, 1)
100000 loops, best of 3: 3.81 µs per loop
In [146]: %timeit d.ravel()[np.arange(d.shape[0])*d.shape[1]+index].reshape(3, 1)
100000 loops, best of 3: 8.56 µs per loop

这种方法的速度是两种方法的两倍。

编辑2:一个更好的方法

基于@Emanuele Paulinis的版本，但减少了操作次数，在大型数组(10k行× 100列)上优于所有操作。

In [199]: %timeit d[(np.arange(index.size), index)].reshape(index.size, 1)
1000 loops, best of 3: 364 µs per loop
In [200]: %timeit d.ravel()[np.arange(d.shape[0])*d.shape[1]+index].reshape(index.size, 1)
100 loops, best of 3: 5.22 ms per loop

如果速度是最重要的:

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