正确索引多维Numpy阵列与另一个索引阵列

我想索引多维数组P与另一个阵列indices。指定哪个沿最后轴我想元，具体如下：正确索引多维Numpy阵列与另一个索引阵列

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]]

来源

2013-03-27 YXD

我可能是缓慢的，但你怎么能（在你的问题“的最后一个轴”）使用二维数组索引沿一个轴？当然，沿着一个轴索引需要一个坐标而不是两个？ – NPE 2013-03-27 14:09:40

我没有解释得很好，这可能与我为什么卡在这里有关。对于任何'i'和'j'我希望'my_output [i，j，：''等于'P [i，j，：，indices [i，j]]' – YXD 2013-03-27 14:13:38

我认为这将工作：

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]

来源

2013-03-27 14:19:39 Jaime

正确索引多维Numpy阵列与另一个索引阵列

回答

相关问题