Numpy: Tri d'un tableau multidimensionnel par un tableau multidimensionnel

Question

Pardonnez-moi si c'est redondant ou super basique. Je viens à Python/Numpy de R et j'ai du mal à faire tourner les choses dans ma tête.

J'ai un tableau de dimension n que je veux trier en utilisant un autre tableau de valeurs n d'indice. Je sais que je pourrais boucler ceci en boucle, mais il semble qu'il devrait y avoir un moyen très concis de le battre en soumission. Voici mon exemple de code pour régler le problème où n = 2:

a1 = random.standard_normal(size=[2,5]) 
index = array([[0,1,2,4,3] , [0,1,2,3,4] ]) 

maintenant j'ai un tableau 2 x 5 de nombres aléatoires et un indice 2 x 5. J'ai lu l'aide pour take() environ 10 fois maintenant mais mon cerveau ne l'entaille pas, évidemment.

Je pensais que cela pourrait me arriver:

take(a1, index) 

array([[ 0.29589188, -0.71279375, -0.18154864, -1.12184984, 0.25698875], 
     [ 0.29589188, -0.71279375, -0.18154864, 0.25698875, -1.12184984]]) 

mais qui est clairement réordonnancement que le premier élément (je suppose en raison de aplanissement).

Des conseils sur comment je vais d'où je suis à une solution qui trie l'élément 0 de a1 par l'élément 0 de l'index ... élément n?

Answer 1

Je ne peux pas penser à la façon de travailler ce dans les dimensions N encore, mais est ici la version 2D:

>>> a = np.random.standard_normal(size=(2,5)) 
>>> a 
array([[ 0.72322499, -0.05376714, -0.28316358, 1.43025844, -0.90814293], 
     [ 0.7459107 , 0.43020728, 0.05411805, -0.32813465, 2.38829386]]) 
>>> i = np.array([[0,1,2,4,3],[0,1,2,3,4]]) 
>>> a[np.arange(a.shape[0])[:,np.newaxis],i] 
array([[ 0.72322499, -0.05376714, -0.28316358, -0.90814293, 1.43025844], 
     [ 0.7459107 , 0.43020728, 0.05411805, -0.32813465, 2.38829386]]) 

Voici la version N-dimensionnelle:

>>> a[list(np.ogrid[[slice(x) for x in a.shape]][:-1])+[i]] 

Voici comment cela fonctionne:

Ok, commençons par un tableau à trois dimensions pour l'illustration.

>>> import numpy as np 
>>> a = np.arange(24).reshape((2,3,4)) 
>>> a 
array([[[ 0, 1, 2, 3], 
     [ 4, 5, 6, 7], 
     [ 8, 9, 10, 11]], 

     [[12, 13, 14, 15], 
     [16, 17, 18, 19], 
     [20, 21, 22, 23]]]) 

Vous pouvez accéder à des éléments de ce tableau en spécifiant l'index le long de chaque axe comme suit:

>>> a[0,1,2] 
6 

Cela équivaut à a[0][1][2] qui est la façon dont vous accéder au même élément si nous avions affaire à une liste au lieu d'un tableau.

Numpy vous permet d'obtenir même colombophile lors du tranchage des tableaux:

>>> a[[0,1],[1,1],[2,2]] 
array([ 6, 18]) 
>>> a[[0,1],[1,2],[2,2]] 
array([ 6, 22]) 

Ces exemples seraient équivalents à [a[0][1][2],a[1][1][2]] et [a[0][1][2],a[1][2][2]] si nous avions affaire à des listes.

Vous pouvez même ignorer des indices répétés et chiffrer le résultat souhaité. Par exemple, les exemples ci-dessus pourraient être écrits de façon équivalente:

>>> a[[0,1],1,2] 
array([ 6, 18]) 
>>> a[[0,1],[1,2],2] 
array([ 6, 22]) 

La forme du tableau (ou liste) vous tranche avec dans chaque dimension affecte uniquement la forme du tableau retourné. En d'autres termes, numpy ne se soucie pas que vous essayez d'indexer votre tableau avec un tableau de forme (2,3,4) quand il tire des valeurs, sauf qu'il va vous renvoyer un tableau de forme (2,3,4).Par exemple:

>>> a[[[0,0],[0,0]],[[0,0],[0,0]],[[0,0],[0,0]]] 
array([[0, 0], 
     [0, 0]]) 

Dans ce cas, nous saisissant le même élément, a[0,0,0] encore et encore, mais numpy retourne un tableau avec la même forme que nous avons passé dans

Ok, sur. ton problème. Ce que vous voulez, c'est indexer le tableau le long du dernier axe avec les numéros de votre tableau index. Donc, pour l'exemple de votre question, vous voulez [[a[0,0],a[0,1],a[0,2],a[0,4],a[0,3]],a[1,0],a[1,1],...

Le fait que votre tableau d'index soit multidimensionnel, comme je l'ai dit plus tôt, ne dit rien sur l'endroit où vous voulez extraire ces indices; il spécifie simplement la forme du tableau de sortie. Donc, dans votre exemple, vous devez indiquer numpy que les 5 premières valeurs doivent être tirées de a[0] et les 5 dernières de a[1]. Facile!

>>> a[[[0]*5,[1]*5],index] 

Il se complique dans les dimensions N, mais faisons-le pour le 3 dimensions a tableau I défini bien au-dessus. Supposons que nous ayons le tableau d'index suivant:

>>> i = np.array(range(4)[::-1]*6).reshape(a.shape) 
>>> i 
array([[[3, 2, 1, 0], 
     [3, 2, 1, 0], 
     [3, 2, 1, 0]], 

     [[3, 2, 1, 0], 
     [3, 2, 1, 0], 
     [3, 2, 1, 0]]]) 

Donc, ces valeurs sont toutes pour les indices le long du dernier axe. Nous devons dire numpy quels indices le long des premier et deuxième axes ces nombres doivent être pris; à-dire que nous devons dire numpy que les indices pour le premier axe sont:

i1 = [[[0, 0, 0, 0], 
     [0, 0, 0, 0], 
     [0, 0, 0, 0]], 

     [[1, 1, 1, 1], 
     [1, 1, 1, 1], 
     [1, 1, 1, 1]]] 

et les indices pour le deuxième axe sont:

i2 = [[[0, 0, 0, 0], 
     [1, 1, 1, 1], 
     [2, 2, 2, 2]], 

     [[0, 0, 0, 0], 
     [1, 1, 1, 1], 
     [2, 2, 2, 2]]] 

alors nous pouvons simplement faire:

>>> a[i1,i2,i] 
array([[[ 3, 2, 1, 0], 
     [ 7, 6, 5, 4], 
     [11, 10, 9, 8]], 

     [[15, 14, 13, 12], 
     [19, 18, 17, 16], 
     [23, 22, 21, 20]]]) 

La fonction numpy pratique qui génère i1 et i2 est appelée np.mgrid. J'utilise np.ogrid dans ma réponse qui est équivalente dans ce cas à cause de la magie numpy dont j'ai parlé plus tôt.

Espérons que ça aide!

Answer 2

Après avoir joué avec ce peu plus aujourd'hui je me suis dit que si j'ai utilisé une fonction mappeur ainsi prendre je pouvais résoudre la version 2 dimensions vraiment simplement comme ceci:

a1 = random.standard_normal(size=[2,5]) 
index = array([[0,1,2,4,3] , [0,1,2,3,4] ]) 
map(take, a1, index) 

je devais map() le take() à Chaque élément dans

Bien sûr, la réponse acceptée résout la version n-dimensionnelle. Cependant, rétrospectivement, j'ai déterminé que je n'ai pas vraiment besoin de la solution n-dimensionnelle, seulement la version 2-D.