Python numpy中的多维索引

假设您有一个

3d numpy数组

，如何沿给定轴构建

N

最大元素的平均值？基本上是这样的：

a = np.random.randint(10, size=(100,100,100)) #axes x,y,z
result = np.mean(a, axis = 2)

但是，我想将平均值限制为

N

沿

轴z

的最大值。为了说明这个问题，这是一个使用循环的解决方案：

a = np.random.randint(10, size=(100,100,100))  #axes x,y,z
N = 5

maxima = np.zeros((100,100,N)) #container for mean of N max values along axis z

for x in range(100):                            #loop through x axis
    for y in range(100):                        #loop through y axis
         max_idx = a[x, y, :].argsort()[-N:]    #indices of N max values along z axis
         maxima[x, y, :] = a[x, y , max_idx]    #extract values

result = np.mean(maxima, axis = 2)              #take the mean

---

我希望通过多维索引获得相同的结果。

这里有一种方法，使用

np.argpartition

获得最大N个索引，然后使用

高级索引

提取和计算所需的平均值-

# Get max N indices along the last axis
maxN_indx = np.argpartition(a,-N, axis=-1)[...,-N:]

# Get a list of indices for use in advanced-indexing into input array,
# alongwith the max N indices along the last axis
all_idx = np.ogrid[tuple(map(slice, a.shape))]
all_idx[-1] = maxN_indx

# Index and get the mean along the last axis
out = a[all_idx].mean(-1)

最后一步也可以用高级索引的显式方式表示，如下所示-

m,n = a.shape[:2]
out = a[np.arange(m)[:,None,None], np.arange(n)[:,None], maxN_indx].mean(-1)