web-dev-qa-db-fra.com

Comment puis-je "trier Zip" des tableaux numpy parallèles?

Si j'ai deux listes parallèles et que je veux les trier selon l'ordre des éléments dans la première, c'est très simple:

>>> a = [2, 3, 1]
>>> b = [4, 6, 2]
>>> a, b = Zip(*sorted(Zip(a,b)))
>>> print a
(1, 2, 3)
>>> print b
(2, 4, 6)

Comment faire la même chose en utilisant des tableaux numpy sans les décompresser dans des listes Python conventionnelles?

36
YGA

b[a.argsort()] devrait faire l'affaire.

Voici comment ça fonctionne. Vous devez d’abord trouver une permutation qui trie un fichier. argsort est une méthode qui calcule ceci:

>>> a = numpy.array([2, 3, 1])
>>> p = a.argsort()
>>> p
[2, 0, 1]

Vous pouvez facilement vérifier que c'est correct:

>>> a[p]
array([1, 2, 3])

Appliquez maintenant la même permutation à b.

>>> b = numpy.array([4, 6, 2])
>>> b[p]
array([2, 4, 6])
48
Jason Orendorff

Voici une approche qui ne crée aucune liste Python intermédiaire, bien qu’elle nécessite un "tableau d’enregistrements" NumPy à utiliser pour le tri. Si vos deux tableaux d'entrée sont réellement liés (comme des colonnes dans une feuille de calcul), cela pourrait constituer un moyen avantageux de traiter vos données en général, plutôt que de conserver deux tableaux distincts tout le temps, auquel cas vous auriez déjà un tableau d’enregistrements et votre problème initial serait résolu simplement en appelant sort () sur votre tableau.

Ceci effectue un tri sur place après avoir placé les deux tableaux dans un tableau d'enregistrement:

>>> from numpy import array, rec
>>> a = array([2, 3, 1])
>>> b = array([4, 6, 2])
>>> c = rec.fromarrays([a, b])
>>> c.sort()
>>> c.f1   # fromarrays adds field names beginning with f0 automatically
array([2, 4, 6])

Edité pour utiliser rec.fromarrays () par souci de simplicité, ignore le type redondant, utilise la clé de tri par défaut, utilise les noms de champs par défaut au lieu de spécifier (basé sur cet exemple ).

20
Peter Hansen

C'est peut-être la façon la plus simple et la plus générale de faire ce que vous voulez. (J'ai utilisé trois tableaux ici, mais cela fonctionnera sur des tableaux de n'importe quelle forme, que ce soit deux colonnes ou deux cents).

import numpy as NP
fnx = lambda : NP.random.randint(0, 10, 6)
a, b, c = fnx(), fnx(), fnx()
abc = NP.column_stack((a, b, c))
keys = (abc[:,0], abc[:,1])          # sort on 2nd column, resolve ties using 1st col
indices = NP.lexsort(keys)        # create index array
ab_sorted = NP.take(abc, indices, axis=0)

Un problème avec lexsort est que vous devez spécifier les clés dans l’ordre inverse, c’est-à-dire que vous devez placer votre clé primaire en premier et votre clé secondaire en premier. Dans mon exemple, je veux trier en utilisant la 2ème colonne comme clé primaire, donc je la liste en second; la 1ère colonne ne résout que les liens, mais elle est listée en premier).

2
doug

Comme la réponse de Peter Hansen, cela crée une copie des tableaux avant de les trier. Mais il est simple: effectue le tri principal sur place, utilise le second tableau pour le tri auxiliaire et devrait être très rapide:

a = np.array([2, 3, 1])
b = np.array([4, 6, 2])
# combine, sort and break apart
a, b = np.sort(np.array([a, b]))
0
Matthias Fripp