Indexing sizes according to number

Question

Problem

I have two arrays numpy, A and indices .

A has dimensions mxnx 10000. indices has dimensions mxnx 5 (output from argpartition(A, 5)[:,:,:5] ). I would like to get an mxnx 5 array containing A elements corresponding to 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]]) A = np.reshape(range(2 * 3 * 10000), (2,3,10000)) A[...,indices] # gives an array of size (2,3,2,3,5). I want a subset of these values np.take(A, indices) # shape is right, but it flattens the array first np.choose(indices, A) # fails because of shape mismatch. 

Motivation

I am trying to get the 5 largest values A[i,j] for each i<m , j<n in sorted order using np.argpartition , because arrays can be quite large.

Answer 1

You can use advanced-indexing -

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

Run Example -

 In [330]: A 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]]]) 

To get these indices corresponding to the maximum 5 elements along the last axis, we will use argpartition , for example:

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

To maintain order from highest to lowest, use range(5) instead of 5 .

Answer 2

For posterity, the following use of Divakar is responsible for achieving the original goal, that is, it returns the top 5 values ​​for all i<m, j<n in sorted order:

 m, n = np.shape(A)[:2] # get the largest 5 indices for all m, n top_unsorted_indices = np.argpartition(A, -5, axis=2)[...,-5:] # get the values corresponding to top_unsorted_indices top_values = A[np.arange(m)[:,None,None], np.arange(n)[:,None], top_unsorted_indices] # sort the top 5 values top_sorted_indices = top_unsorted_indices[np.arange(m)[:,None,None], np.arange(n)[:,None], np.argsort(-top_values)]