使用另一个索引数组正确地对多维Numpy数组进行索引

Question

我正在尝试使用另一个数组indices来索引一个多维数组P。该数组指定我想要沿着最后一个轴选择哪个元素，如下所示：

import numpy as np

M, N = 20, 10

P = np.random.rand(M,N,2,9)

# index into the last dimension of P
indices = np.random.randint(0,9,size=(M,N))

# I'm after an array of shape (20,10,2)
# but this has shape (20, 10, 2, 20, 10)
P[...,indices].shape 

如何使用 indices 正确地索引 P，以获得形状为 (20,10,2) 的数组?

如果不太清楚：对于任何在范围内的 i 和 j，我希望 my_output[i,j,:] 等于 P[i,j,:,indices[i,j]]

- YXD

---

我可能有点慢，但是你如何使用2D数组沿着一个轴（在你的问题中是“最后一个轴”）进行索引呢？毕竟，沿着单个轴进行索引只需要一个坐标而不是两个吧？ - NPE

我没有解释得很好，这可能与我现在卡住了有关。对于任何 i 和 j，我希望 my_output[i,j,:] 等于 P[i,j,:,indices[i,j]]。 - YXD

Answer 1

我认为这个方案可行：

P[np.arange(M)[:, None, None], np.arange(N)[:, None], np.arange(2),
  indices[..., None]]

我知道这不太美观...

---

这个看起来可能更好看，但可读性可能会降低：

P[np.ogrid[0:M, 0:N, 0:2]+[indices[..., None]]]

或者更好的选择：

idx_tuple = tuple(np.ogrid[:M, :N, :2]) + (indices[..., None],)
P[idx_tuple]

- Jaime