python - Numpy:按数组索引数组

Question

>>> idx = np.random.randint(2, size=(9, 31))
>>> a = np.random.random((9, 31, 2))
>>> a[idx].shape
(9, 31, 31, 2)

为什么上面没有产生至少 (9, 31, 1) 或更好的 (9, 31) 形状？我怎样才能让它根据 idx 中的值返回一个选择？

更新

这可能是一个更具体且更希望类似的示例:假设这个数组

a = np.asarray([[1, 2], [3, 4], [5, 6], [7, 8]])

我将如何选择数组 [1, 4, 5, 8](即每行的第 0、1、0、1 个元素)？

Answer 1

我想这就是你想要的:

>>> a[np.arange(9)[:, None], np.arange(31), idx].shape
(9, 31)

对于你的第二个例子，你会这样做:

>>> a[np.arange(4), [0, 1, 0, 1]]
array([1, 4, 5, 8])

阅读the docs关于花式索引，尤其是当你没有每个维度的索引数组时会发生什么的部分 here : 那些额外的 np.arange 数组被放置在那里以避免这种行为。

还要注意它们是如何 reshape 的(使用 [:, None] 索引等同于 .reshape(-1, 1))以便它们的广播形状具有所需输出数组的形状。