numpyでの多次元インデクシング

Question

3d numpy arrayがあるとします。Nの要素の平均を特定の軸に沿ってどのように構築できますか？だから、基本的には何か：

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

は、しかし、私はaxis z沿っN最大値に平均値を制限したいです。この問題を説明するために、これはループを使用するソリューションです。

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 

多次元索引付けで同じ結果を得たいと思います。

Answer 1

ここで所望の平均値を抽出して計算するための1つの最大N個のインデックスを取得するnp.argpartitionを用いてアプローチした後advanced-indexingだ -

# 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) 

最後のステップも同様に、高度、索引付けのための明示的な方法で表現することができるがso -

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