访问numpy数组中行特定元素的矢量化方法 [英] Vectorized way of accessing row specific elements in a numpy array

问题描述

我有一个二维NumPy数组和一组索引，其大小是NumPy数组的第一维.

I have a 2-D NumPy array and a set of indices the size of which is the first dimension of the NumPy array.

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

我需要做类似的事情

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

有矢量化的方法吗?

推荐答案

在这里:

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

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

I'm assuming here that you don't need each value on a separate line and just want to extract the values.

这篇关于访问numpy数组中行特定元素的矢量化方法的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！