用另一个索引数组正确索引多维Numpy数组 [英] Correctly indexing a multidimensional Numpy array with another array of indices
本文介绍了用另一个索引数组正确索引多维Numpy数组的处理方法,对大家解决问题具有一定的参考价值,需要的朋友们下面随着小编来一起学习吧!
问题描述
我正在尝试用另一个数组indices
为多维数组P
编制索引.它指定了我想要的最后一个轴上的哪个元素,如下所示:
I'm trying to index a multidimensional array P
with another array indices
. which specifies which element along the last axis I want, as follows:
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)
的数组?
How can I correctly index P
with indices
to get an array of shape (20,10,2)
?
如果不太清楚:对于任何i
和j
(在边界内),我希望my_output[i,j,:]
等于P[i,j,:,indices[i,j]]
If that's not too clear: For any i
and j
(in bounds) I want my_output[i,j,:]
to be equal to P[i,j,:,indices[i,j]]
推荐答案
我认为这会起作用:
P[np.arange(M)[:, None, None], np.arange(N)[:, None], np.arange(2),
indices[..., None]]
不漂亮,我知道...
Not pretty, I know...
这可能看起来更好,但也可能不太清晰:
This may look nicer, but it may also be less legible:
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数组的文章就介绍到这了,希望我们推荐的答案对大家有所帮助,也希望大家多多支持IT屋!
查看全文