python - Numpy match indexing dimensions -

problem

i have 2 numpy arrays, a , indices.

a has dimensions m x n x 10000. indices has dimensions m x n x 5 (output argpartition(a, 5)[:,:,:5]). m x n x 5 array containing elements of a corresponding indices.

attempts

indices = np.array([[[5,4,3,2,1],[1,1,1,1,1],[1,1,1,1,1]],     [500,400,300,200,100],[100,100,100,100,100],[100,100,100,100,100]]) = np.reshape(range(2 * 3 * 10000), (2,3,10000))  a[...,indices] # gives array of size (2,3,2,3,5). want subset of these values np.take(a, indices) # shape right, flattens array first np.choose(indices, a) # fails because of shape mismatch.  

motivation

i'm trying 5 largest values of a[i,j] each i<m, j<n in sorted order using np.argpartition because arrays can large.

you can use advanced-indexing -

m,n = a.shape[:2] out = a[np.arange(m)[:,none,none],np.arange(n)[:,none],indices] 

sample run -

in [330]: out[330]:  array([[[38, 21, 61, 74, 35, 29, 44, 46, 43, 38],         [22, 44, 89, 48, 97, 75, 50, 16, 28, 78],         [72, 90, 48, 88, 64, 30, 62, 89, 46, 20]],         [[81, 57, 18, 71, 43, 40, 57, 14, 89, 15],         [93, 47, 17, 24, 22, 87, 34, 29, 66, 20],         [95, 27, 76, 85, 52, 89, 69, 92, 14, 13]]])  in [331]: indices out[331]:  array([[[7, 8, 1],         [7, 4, 7],         [4, 8, 4]],         [[0, 7, 4],         [5, 3, 1],         [1, 4, 0]]])  in [332]: m,n = a.shape[:2]  in [333]: a[np.arange(m)[:,none,none],np.arange(n)[:,none],indices] out[333]:  array([[[46, 43, 21],         [16, 97, 16],         [64, 46, 64]],         [[81, 14, 43],         [87, 24, 47],         [27, 52, 95]]]) 

for getting indices corresponding max 5 elements along last axis, use argpartition, -

indices = np.argpartition(-a,5,axis=-1)[...,:5] 

to keep order highest lowest, use range(5) instead of 5.