用另一个数组索引一个numpy数组[重复]

Question

我有一个 numpy 数组

a = np.random.random((10,3))

我想使用b 来索引a

# array([0, 2, 1, 2, 0, 0, 2, 0, 0, 2])
b = np.random.choice(a=a.shape[1], size=a.shape[0])

我希望结果是(10, )，但我尝试了以下方法，但都失败了

print(a[b].shape) # (10, 3)
print(a[b,:].shape) # (10, 3)
print(a[:,b].shape) # (10, 10)
print(a[b.reshape(10,1)].shape) # (10, 1, 3)
print(a[b.reshape(1,10)].shape) # (1, 10, 3)

在我的情况下，索引a 的正确方法是什么？

Answer 1

尝试使用np.arange(a.shape[0]) 作为行索引，使用数组b 作为列索引：

>>> a[np.arange(a.shape[0]), b].shape
(10,)