Основные способы сортировки массивов в NumPy

• 29.08.2023
• 744
• Библиотеки Python
• NumPy

Сортировка — одна из базовых и важных операций в NumPy. В библиотеке имеется несколько функций и методов для выполнения сортировки массивов.

Функция sort()

Эта функция возвращает отсортированную копию массива.

import numpy as np

arr = np.array([3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5])
sorted_arr = np.sort(arr)
print(sorted_arr)  # [1 1 2 3 3 4 5 5 5 6 9]

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

import numpy as np

matrix = np.array([[34, 23, 19], [19, 87, 56], [81, 42, 64]])
sorted_matrix = np.sort(matrix, axis=1)
print(sorted_matrix)
# [[19 23 34]
#  [19 56 87]
#  [42 64 81]]

Метод ndarray.sort()

Этот метод сортирует массив на месте, то есть изменяет исходный массив:

import numpy as np

arr = np.array([3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5])
arr.sort()
print(arr)  # [1 1 2 3 3 4 5 5 5 6 9]

Функция argsort()

Возвращает индексы, которые будут использоваться для сортировки массива:

import numpy as np

arr = np.array([3, 1, 4, 1])
indices = np.argsort(arr)
print(indices)  # [1 3 0 2]

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

Функция lexsort()

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

import numpy as np

a = [1, 5, 1]
b = [9, 4, 0]
ind = np.lexsort((a, b))  # Сортировка сначала по b, затем по a
print(ind)  # [2 1 0]

Сортировка структурированных массивов

NumPy позволяет создавать структурированные массивы, которые можно сортировать по определенным полям:

import numpy as np

structured_arr = np.array([('Rex', 9, 81.0), ('Fido', 3, 27.0)],
                          dtype=[('name', 'U10'), ('age', 'i4'), ('weight', 'f4')])
print(np.sort(structured_arr, order='age'))

# [('Fido', 3, 27.) ('Rex', 9, 81.)]

Функция partition()

Эта функция возвращает массив, где первые k элементов — это наименьшие k элементов в произвольном порядке, а остальные элементы размещены в произвольном порядке за ними:

import numpy as np

arr = np.array([3, 1, 4, 1, 5, 9, 2, 6, 5])
print(np.partition(arr, 3))  # [1 1 2 3 4 5 9 6 5]

Функция argpartition()

Возвращает индексы, которые бы частично отсортировали массив так, чтобы элемент на k-той позиции был на своем отсортированном месте. Элементы перед k-тым будут меньше, а после — больше, но они не обязательно будут в отсортированном порядке.

import numpy as np

arr = np.array([3, 1, 4, 1, 5, 9, 2, 6, 5])
indices = np.argpartition(arr, 3)
print(indices)  # [3 1 6 0 2 4 5 7 8]

Заключение

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

Тем не менее, выбор правильного метода сортировки требует понимания их особенностей и различий. Например, использование argsort() может быть крайне полезным, когда необходимо сохранить индексы исходного массива, в то время как partition() подойдет для быстрого поиска топ-N элементов без полной сортировки.