Если у меня есть два параллельных списка и я хочу отсортировать их по порядку элементов в первом, это очень просто:
>>> 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?
b[a.argsort()] должен сделать b[a.argsort()] дело.
Вот как это работает. Сначала вам нужно найти перестановку, которая сортирует. argsort это метод, который вычисляет это:
>>> a = numpy.array([2, 3, 1])
>>> p = a.argsort()
>>> p
[2, 0, 1]
Вы можете легко проверить, что это правильно:
>>> a[p]
array([1, 2, 3])
Теперь примените ту же перестановку к б.
>>> b = numpy.array([4, 6, 7])
>>> b[p]
array([7, 4, 6])
Здесь подход, который не создает промежуточных списков 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, использования ключа сортировки по умолчанию, использования имен полей по умолчанию вместо указания (на основе этого примера).
Это может быть самый простой и самый общий способ делать то, что вы хотите. (Здесь я использовал три массива, но это будет работать с массивами любой формы, будь то две колонки или две сотни).
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)
Один quirk w/lexsort заключается в том, что вы должны указать ключи в обратном порядке, то есть сначала поместить свой первичный ключ и ваш вторичный ключ. В моем примере я хочу сортировать, используя второй столбец в качестве первичного ключа, поэтому я перечислил его вторым; 1-й столбец разрешает только ссылки, но он указан первым).
Как и ответ @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 была написана 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]))