NumPy mecz wymiary indeksujące

Question

Problem

Mam dwa NumPy tablic A i indices.

A ma wymiary M x N x 10000 indices ma wymiary M x N x 5 (wyjściowy argpartition(A, 5)[:,:,:5]). Chciałbym uzyskać tablicę m x n x 5 zawierającą elementy A odpowiadającą indices.

Próby

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. 

Motywacja

Próbuję dostać 5 największych wartości A[i,j] dla każdego i<m, j<n w kolejności sortowane przy użyciu np.argpartition ponieważ tablice można uzyskać dość duże.

Answer 1

Można użyć advanced-indexing -

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

run Sample -

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

Na uzyskanie tych wskaźników odpowiadających maksymalnie 5 elementów wzdłuż ostatniej osi, to użyłby argpartition, jak tak -

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

Aby utrzymać zamówienie od najwyższej do najniższej, nas e range(5) zamiast 5.

Answer 2

dla potomnych, co następuje wykorzystuje odpowiedź Divakar do osiągnięcia pierwotnego celu, czyli powrót top 5 wartości dla wszystkich i<m, j<n posortowanych:

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