在numpy数组中访问行特定元素的矢量化方式

Question

我有一个2-D NumPy数组和一组索引，其大小是NumPy数组的第一个维度。

X = np.random.rand(5, 3)
a = np.random.randint(0, 3, 5)

我需要做类似

的事情

for i, ind in enumerate(a):
    print X[i][ind]

有没有矢量化的方法呢？

Answer 1

你走了：

X = np.random.rand(5, 3)
a = np.random.randint(0, 3, 5)

In [12]: X[np.arange(a.size), a]
Out[12]: array([ 0.99653335,  0.30275346,  0.92844957,  0.54728781,  0.43535668])
In [13]: for i, ind in enumerate(a):
            print X[i][ind]
#   ....:
#0.996533345844
#0.30275345582
#0.92844956619
#0.54728781105
#0.435356681672

我假设您不需要在单独的一行上提供每个值，只想提取值。