Multidimensional indexing in numpy

Question

Suppose you have a 3d numpy array , how can I build the mean of the N maximum elements along a given axis? So basically something like:

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

however, I want to restrict the mean to the N maximum values along axis z . To illustrate the issue, this is a solution using looping:

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

I would like to achieve the same result with multidimensional indexing.

Answer 1

Here's one approach using np.argpartition to get the max N indices and then advanced-indexing for extracting and computing the desired average values -

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

Last step could also be expressed in an explicit way for advanced-indexing, like so -

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