从数组中选择每行中的特定列

Question

我试图为numpy数组的每一行选择特定的列元素.例如,在以下示例中:

In [1]: a = np.random.random((3,2))
Out[1]: 
array([[ 0.75670668,  0.1283942 ],
       [ 0.51326555,  0.59378083],
       [ 0.03219789,  0.53612603]])

我想选择第一行的第一个元素,第二行的第二个元素和第三行的第一个元素.所以我尝试做以下事情:

In [2]: b = np.array([0,1,0])

In [3]: a[:,b]

但是这会产生以下输出:

Out[3]: 
array([[ 0.75670668,  0.1283942 ,  0.75670668],
       [ 0.51326555,  0.59378083,  0.51326555],
       [ 0.03219789,  0.53612603,  0.03219789]])

这显然不是我想要的.有没有一种简单的方法可以在不使用循环的情况下做我想做的事情？

Answer 1

您可以使用:

a[np.arange(3), (0,1,0)]

在上面的例子中.