Получение индексов отсортированного списка в Python

Быстрый ответ

Для получения индексов отсортированного массива в Python можно воспользоваться функцией numpy.argsort():

import numpy as np

array = [3, 1, 2]
indices = np.argsort(array)

print(indices)  # Результат: [1, 2, 0]

Функция np.argsort() возвращает индексы, с помощью которых массив array можно упорядочить в порядке возрастания. В исходном массиве [3, 1, 2] после сортировки порядок элементов изменится на [1, 2, 3], что соответствует индексам [1, 2, 0].

Что делать, если нет NumPy?

Если вы не установили NumPy, вы всегда можете воспользоваться стандартными возможностями Python. Для отсортировки пар (индекс, значение) по значению примените выражение sorted(enumerate(list), key=lambda x: x[1]):

array = [5, 1, 3]
sorted_indices = [i for (i, _) in sorted(enumerate(array), key=lambda x: x[1])]

print(sorted_indices)  # Выходные данные: [1, 2, 0]

Если вам не по душе использование функции enumerate, тогда sorted(range(len(list)), key=list.__getitem__) подскажет более простой способ получения индексов:

array = [5, 1, 3]
sorted_indices = sorted(range(len(array)), key=array.__getitem__)

print(sorted_indices)  # Результат: [1, 2, 0]

Продвинутые техники сортировки, сохраняющие исходный порядок

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

array = [3, 1, 2, 1]
sorted_indices = sorted(range(len(array)), key=lambda i: (array[i], i))

print(sorted_indices)  # Сохраняется первоначальный порядок '1', результат: [1, 3, 2, 0]

При сложных сравнениях или работы с пользовательскими объектами, можно представить параметр key в виде собственной функции.

Визуализация

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

Исходное расположение на полке: [📗, 📘, 📙, 📕]
- 📗: Зелёная книга
- 📘: Синяя книга
- 📙: Оранжевая книга
- 📕: Красная книга

Чтобы не потерять исходное место каждой книги после сортировки, вы используете закладки с номерами:

sorted_indices = [3, 0, 1, 2]

Новое расположение на полке: [📕, 📗, 📘, 📙]
Номера закладок:              [#4, #1, #2, #3]

Так выглядит визуализация процесса сортировки индексов массива в Python — номера закладок подобны индексам в вашем массиве!

Более сложные сценарии сортировки

Убывающий порядок

Если требуется отсортировать индексы в убывающем порядке, достаточно добавить минус (-) перед каждым элементом при сортировке:

array = [3, 1, 2]
sorted_indices = np.argsort([-a for a in array])

print(sorted_indices)  # Результат: [0, 2, 1]

Работа с многомерными массивами

При работе с многомерными массивами или матрицами, можно указать параметр axis, чтобы управлять сортировкой по строкам или столбцам:

matrix = np.array([[3, 2], [1, 0]])
row_sort_indices = np.argsort(matrix, axis=1)
column_sort_indices = np.argsort(matrix, axis=0)

print(row_sort_indices)      # Результат: [[1, 0], [1, 0]]
print(column_sort_indices)   # Результат: [[1, 1], [0, 0]]

Работа с нетиповыми данными

Для сортировки более сложных структур данных и объектов, методы attrgetter или itemgetter из модуля operator могут быть весьма полезны:

from operator import itemgetter

pairs = [(2, 'apple'), (1, 'banana'), (2, 'cherry')]
sorted_indices = sorted(range(len(pairs)), key=itemgetter(0))

print(sorted_indices)  # Банан возглавляет список! Результат: [1, 0, 2]

Производительность — наше всё

Оцениваем производительность с помощью perfplot

Для оценки производительности различных методов сортировки можно использовать инструменты, такие как perfplot. С помощью наглядных графиков вы сможете сравнить производительность np.argsort() и других подходов.

Работа с большими наборами данных

При работе с большими массивами данных важно учесть алгоритмическую сложность. Функция NumPy argsort оптимизирована для работы с большими объемами данных, но стандартные Python-функции в совокупности с списочными включениями могут быть менее эффективными в подобных условиях.

Полезные материалы

1. numpy.argsort — Руководство NumPy v1.26: для дополнительной информации о функции argsort обратитесь к официальной документации NumPy.
2. Встроенные функции — Документация Python 3.12.2: хотите узнать больше о возможностях встроенных функций Python, в том числе функции sorted()?
3. 5. Структуры данных — Документация Python 3.12.2: окунитесь в мир списочных включений, одного из важных инструментов Python.
4. Как сортировать — Документация Python 3.12.2: полезное руководство по вопросам сортировки и синтаксису ключевых функций.
5. pandas.DataFrame.sort_values — Документация pandas 2.2.0: разъяснение принципов сортировки значений DataFrame в официальной документации pandas.
6. Python enumerate(): Упрощенные циклы с индексами – Real Python: подробная статья от Real Python о функции enumerate(), которая поможет разобраться в деталях ее применения.