Question

我正在尝试使用另一个数组P索引多维数组indices。它指定了我想要的最后一个轴上的哪个元素，如下所示：

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

如何使用P正确索引indices以获取形状(20,10,2)的数组？

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

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]

使用另一个索引数组正确索引多维Numpy数组

1 个答案: