Как я могу "отсортировать по zip" параллельные массивы numpy?

Question

Если у меня есть два параллельных списка и я хочу отсортировать их по порядку элементов в первом, это очень просто:

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

Как я могу сделать то же самое, используя массивы numpy, не распаковывая их в обычные списки Python?

Answer 1

b[a.argsort()] должен сделать трюк.

Вот как это работает. Сначала вам нужно найти перестановку, которая сортирует. argsort это метод, который вычисляет это:

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

Вы можете легко проверить, что это правильно:

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

Теперь примените ту же перестановку к b.

>>> b = numpy.array([4, 6, 7])
>>> b[p]
array([7, 4, 6])

Answer 2

Вот подход, который не создает промежуточных списков Python, хотя для сортировки ему необходим «массив записей» NumPy. Если ваши два входных массива на самом деле связаны (например, столбцы в электронной таблице), это может открыть полезный способ работы с вашими данными в целом, вместо того, чтобы хранить два разных массива постоянно, в этом случае у вас уже есть массив записей, и ваша первоначальная проблема будет решена простым вызовом sort () вашего массива.

Это делает сортировку на месте после упаковки обоих массивов в массив записей:

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

Изменено для простоты использования rec.fromarrays (), пропуска избыточного dtype, использования ключа сортировки по умолчанию, использования имен полей по умолчанию вместо указания (на основе в этом примере ).

Answer 3

Как и ответ @Peter Hansen, он создает копии массивов, а затем сортирует их. Но он прост, выполняет основную сортировку на месте, использует второй массив для вспомогательной сортировки и должен быть очень быстрым:

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

Обновление : приведенный выше код фактически не работает, как указано в комментарии. Ниже приведен лучший код. Это должно быть достаточно эффективным, например, избегать явного создания дополнительных копий массивов. Трудно сказать, насколько это будет эффективно, потому что документация не дает никаких подробностей об алгоритме numpy.lexsort. Но это должно работать очень хорошо, так как это именно та работа, для которой было написано lexsort.

a = np.array([5, 3, 1])
b = np.array([4, 6, 7])
new_order = np.lexsort([b, a])
a = a[new_order]
b = b[new_order]
print(a, b)
# (array([1, 3, 5]), array([7, 6, 4]))

Answer 4

Это может быть самый простой и общий способ сделать то, что вы хотите. (Я использовал три массива здесь, но это будет работать с массивами любой формы, будь то два столбца или двести).

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)

Одна из странностей с lexsort заключается в том, что вы должны указывать ключи в обратном порядке, то есть ставить первичный ключ вторым, а вторичный ключ первым. В моем примере я хочу отсортировать, используя 2-й столбец в качестве первичного ключа, поэтому я перечислю его вторым; 1-й столбец разрешает только связи, но он указан первым).