Renvoyer les N plus grands indices de valeurs dans un tableau multidimensionnel (peut trouver des solutions pour une dimension mais pas multi-dimension)

Question

J'ai un tableau num X, et je voudrais retourner un autre tableau Y dont les entrées sont les indices de les n valeurs les plus élevées de X-à-dire suppose que j'ai:

a =np.array[[1, 3, 5], [4, 5 ,6], [9, 1, 7]] 

puis dire, si je veux que les 5 premiers "maxs" indices-ici 9 s ', 7, 6, 5, 5 sont les maxs, et leur indices sont les suivants:

b=np.array[[2, 0], [2 2], [ 2 1], [1 1], [0 , 2]) 

Je suis en mesure de trouver des solutions et faire ce travail pour un tableau à deux dimensions comme

c=np.array[1, 2, 3, 4, 5, 6]: 

def f(a,N): 
    return np.argsort(a)[::-1][:N] 

Mais n'ont pas été en mesure de générer quelque chose qui fonctionne dans plus d'une dimension. Merci!

Answer 1

Approche # 1

Obtenez les indices argsort sur sa version aplatis et sélectionnez les derniers N indices. Ensuite, obtenir les indices de ligne et colonne correspondant -

N = 5 
idx = np.argsort(a.ravel())[-N:][::-1] #single slicing: `[:N-2:-1]` 
topN_val = a.ravel()[idx] 
row_col = np.c_[np.unravel_index(idx, a.shape)] 

run échantillon -

# Input array 
In [39]: a = np.array([[1,3,5],[4,5,6],[9,1,7]]) 

In [40]: N = 5 
    ...: idx = np.argsort(a.ravel())[-N:][::-1] 
    ...: topN_val = a.ravel()[idx] 
    ...: row_col = np.c_[np.unravel_index(idx, a.shape)] 
    ...: 

In [41]: topN_val 
Out[41]: array([9, 7, 6, 5, 5]) 

In [42]: row_col 
Out[42]: 
array([[2, 0], 
     [2, 2], 
     [1, 2], 
     [1, 1], 
     [0, 2]]) 

---

Approche # 2

Pour des performances, nous pouvons utiliser np.argpartition pour obtenir top N indices sans garder l'ordre trié, comme si -

idx0 = np.argpartition(a.ravel(), -N)[-N:] 

Pour obtenir l'ordre de tri, nous avons besoin d'une autre série de argsort -

idx = idx0[a.ravel()[idx0].argsort()][::-1]