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

Question

我有一个二维 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

我在这里假设您不需要将每个值放在单独的行上，而只想提取这些值。