我试图用另一个数组索引索引多维数组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
如何使用索引正确地索引P以获得形状数组(20,10,2)?
如果那不太清楚:对于任何i和j(在界限中)我希望my_output [i,j,:]等于P [i,j,:,indices [i,j]]
最佳答案 我认为这会奏效:
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]