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

2024-05-05

我正在尝试索引多维数组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 with indices获得形状数组(20,10,2)?

如果这还不太清楚:对于任何i and 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]
本文内容由网友自发贡献,版权归原作者所有,本站不承担相应法律责任。如您发现有涉嫌抄袭侵权的内容,请联系:hwhale#tublm.com(使用前将#替换为@)

使用另一个索引数组正确索引多维 Numpy 数组 的相关文章

随机推荐