numpy.lexsort

numpy.lexsort(keys, axis=-1)

Функция lexsort() выполняет обратную устойчивую сортировку на основе указанных ключей (лексикографическая сортировка).

Параметры:
keys - двумерный массив NumPy, кортеж одномерных массивов или кортеж массивоподобных объектов.
Сортируемые одномерные последовательности или массивы. Если это двумерный массив, то все его строки рассматриваются как, отдельные одномерные последовательности, а их порядок в нем, как порядок ключей. Последний ключ или строка двумерного массива сортируется в первую очередь.
axis - целое число (необязательный параметр).
Определяет ось вдоль которой выполняется сортировка элементов. По умолчанию axis = -1, что соответствует сортировке по последней оси массива.
Возвращает:
ndarray - массив NumPy
Одномерный массив индексов, который сортирует массивы в указанном в keys порядке и вдоль указанной оси.
Смотрите так же:
argsort, sort


Примеры

Данная функция выполняет сортировку над массивами, как над единым целым, а последовательность ключей указывает лишь на порядок сортировки. Заметьте, что в первую очередь сортировка выполняется по ключу, который указан последним:

>>> import numpy as np
>>> 
>>> A = ['c', 'a', 'a', 'b']
>>> B = ['dd', 'cc', 'bb', 'aa']
>>> 
>>> 
>>> ind_a = np.lexsort((B, A))    #  сначала по 'A'
>>> ind_a
array([2, 1, 3, 0], dtype=int32)
>>> 
>>> [(A[i], B[i]) for i in ind_a]
[('a', 'bb'), ('a', 'cc'), ('b', 'aa'), ('c', 'dd')]
>>> 
>>> 
>>> ind_b = np.lexsort((A, B))    #  сначала по 'B'
>>> ind_b
array([3, 2, 1, 0], dtype=int32)
>>> 
>>> [(B[i], A[i]) for i in ind_b]
[('aa', 'b'), ('bb', 'a'), ('cc', 'a'), ('dd', 'c')]

По сути, все одномерные последовательности рассматриваются как единое целое, а сортировка выполняется не по отдельным массивам, а по их совокупности. Например, три отдельных массива, будут рассматриваться как один массив, а возвращаться будут его отсортированные столбцы:

>>> a = [0, 1, 1, 0, 1, 0]
>>> b = [99, 17, 46, 53, 14, 87]
>>> c = [834, 567, 997, 112, 342, 566]
>>> 
>>> ind = np.lexsort((c, b, a))
>>> ind
array([3, 5, 0, 4, 1, 2], dtype=int32)
>>> 
>>> 
>>> for i in ind:
...     print(a[i], b[i], c[i])
... 
0 53 112
0 87 566
0 99 834
1 14 342
1 17 567
1 46 997

Пример выполненный выше, равносилен сортировке аналогичного двумерного массива:

>>> a = [0, 1, 1, 0, 1, 0]
>>> b = [99, 17, 46, 53, 14, 87]
>>> c = [834, 567, 997, 112, 342, 566]
>>> 
>>> arr  =np.array([c, b, a])
>>> arr
array([[834, 567, 997, 112, 342, 566],
       [ 99,  17,  46,  53,  14,  87],
       [  0,   1,   1,   0,   1,   0]])
>>> 
>>> ind = np.lexsort(arr)
>>> ind
array([3, 5, 0, 4, 1, 2], dtype=int32)
>>> 
>>> arr[:, ind]
array([[112, 566, 834, 342, 567, 997],
       [ 53,  87,  99,  14,  17,  46],
       [  0,   0,   0,   1,   1,   1]])

Что бы явно увидеть отличие сортировки по всем ключам от сортировки по дному ключу, достаточно воспользоваться функцией np.argsort():

>>> arr[:, np.argsort(a)]
array([[834, 112, 566, 567, 997, 342],
       [ 99,  53,  87,  17,  46,  14],
       [  0,   0,   0,   1,   1,   1]])

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

>>> a = np.array([[1, 1], [2, 2]])
>>> b = np.array([[3, 3], [4, 4]])
>>> 
>>> np.lexsort((b, a))
array([[0, 1],
       [0, 1]], dtype=int32)
>>> 
>>> np.lexsort((b, a), axis = 0)
array([[0, 0],
       [1, 1]], dtype=int32)

Кстати, выполнить лексикографическую сортировку структурированного массива можно с помощью той же функции - np.argsort():

>>> arr = np.array([(0, 12), (0, 11), (1, 55),
...                 (0, 87), (1, 54), (1, 21)],
...                dtype = np.dtype([('a', int), ('b', int)]))
>>> arr
array([(0, 12), (0, 11), (1, 55), (0, 87), (1, 54), (1, 21)],
      dtype=[('a', '<i4'), ('b', '<i4')])
>>> 
>>> ind = np.argsort(arr, order = ['a', 'b'])
>>> ind
array([1, 0, 3, 5, 4, 2], dtype=int32)
>>> 
>>> arr[ind]
array([(0, 11), (0, 12), (0, 87), (1, 21), (1, 54), (1, 55)],
      dtype=[('a', '<i4'), ('b', '<i4')])
>>> 
>>> 
>>> ind = np.argsort(arr, order = ['b', 'a'])
>>> ind
array([1, 0, 5, 4, 2, 3], dtype=int32)
>>> 
>>> arr[ind]
array([(0, 11), (0, 12), (1, 21), (1, 54), (1, 55), (0, 87)],
      dtype=[('a', '<i4'), ('b', '<i4')])