是否可以使用numpy.take来选择索引列表

Question

我想从numpy.array的元素中获取一些索引

import numpy as np
a = np.array(range(20))
a.shape=4,5
#~ [[ 0  1  2  3  4]
#~ [ 5  6  7  8  9]
#~ [10 11 12 13 14]
#~ [15 16 17 18 19]]
idxs = np.array([(4,0),(3,1),(2,1),(0,3)])
ret = a.take(idxs,1).diagonal().transpose()
print ret
#~ [[  4.   0.]
#~ [  8.   6.]
#~ [ 12.  11.]
#~ [ 15.  18.]]

有可能以更简单（或更快）的方式得到这个结果吗？

Answer 1

您可以为轴0索引创建一个数组：

a[np.arange(4)[:,None], idxs]