Question

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

例如

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

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

更新

由于我会有更大的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，我和＃ 39; ll采取：）

Answer 1

可能有更好的方法，但take，diag和reshape的组合可以做到：

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

修改

与@Emanuele Paolinis＆＃39;比较改变，添加重塑以匹配所寻求的输出：

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＆＃39;版本但减少的操作数量优于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)

Answer 2

这可能很难看但效率更高：

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

修改

正如@deinonychusaur所指出的，上面的陈述可以写成干净的：

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

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

2 个答案: