Функция lexsort

Функция lexsort выполняет непрямую сортировку по нескольким ключам. В качестве первого аргумента функция принимает последовательность сортируемых массивов или кортеж массивов. Важно: сортировка происходит по последнему переданному ключу в первую очередь. Вторым параметром можно передать ось, по которой производится сортировка. Функция возвращает массив индексов, которые могут быть использованы для сортировки исходного массива данных.

Синтаксис

np.lexsort(keys, [axis])

Пример с одним ключом

Давайте отсортируем массив с помощью одного ключа:

import numpy as np

arr = np.array([3, 1, 4, 2])
res = np.lexsort((arr,))
print(res)

Результат выполнения кода:

[1 3 0 2]

Пример сортировки по нескольким ключам

Отсортируем строки таблицы по двум столбцам. Сначала сортировка по последнему ключу, затем по первому:

import numpy as np

first = np.array([2, 1, 2, 1])
second = np.array([3, 2, 1, 4])
res = np.lexsort((first, second))
print(res)

Результат выполнения кода:

[2 1 0 3]

Пример с двумерным массивом

Сортируем строки матрицы по первому столбцу, а затем по второму:

import numpy as np

arr = np.array([[3, 2], [1, 4], [3, 1], [1, 3]])
res = np.lexsort((arr[:, 1], arr[:, 0]))
print(res)

Результат выполнения кода:

[1 3 2 0]

Пример применения индексов к массиву

Используем полученные индексы для сортировки исходного массива:

import numpy as np

names = np.array(['Anna', 'Bob', 'Anna', 'Bob'])
grades = np.array([85, 92, 78, 88])
res = np.lexsort((grades, names))
sorted_names = names[res]
sorted_grades = grades[res]
print(sorted_names)
print(sorted_grades)

Результат выполнения кода:

['Anna' 'Anna' 'Bob' 'Bob']
[78 85 88 92]

Пример с отрицательной сортировкой

Для сортировки по убыванию используем отрицательные значения ключа:

import numpy as np

first = np.array([2, 1, 2, 1])
second = np.array([3, 2, 1, 4])
res = np.lexsort((-first, second))
print(res)

Результат выполнения кода:

[0 2 1 3]

Смотрите также

функцию sort,
которая сортирует массив на месте
функцию argsort,
которая выполняет непрямую сортировку по одному ключу
функцию partition,
которая выполняет частичную сортировку массива
функцию searchsorted,
которая выполняет поиск в отсортированном массиве