Метод argpartition
Метод argpartition класса ndarray возвращает индексы элементов, которые заняли бы свои места при частичной сортировке массива. Первым параметром метод принимает индекс kth, который определяет позицию, на которую должен встать элемент. Вторым параметром можно задать ось для многомерных массивов. Метод полезен, когда нужно найти позиции нескольких наименьших или наибольших элементов.
Синтаксис
arr.argpartition(kth, axis=-1, kind='introselect', order=None)
Пример
Давайте найдём индексы двух наименьших элементов в одномерном массиве:
import numpy as np
np.random.seed(0)
arr = np.random.randint(0, 100, 10)
print("Original array:", arr)
res = arr.argpartition(1)
print("Indices:", res)
print("Two smallest values:", arr[res[:2]])
Результат выполнения кода:
"Original array: [44 47 64 67 67 9 83 21 36 87]"
"Indices: [5 7 8 3 2 1 9 0 4 6]"
"Two smallest values: [ 9 21]"
Пример
Найдём индексы трёх наибольших элементов в массиве. Для этого используем отрицательное значение kth:
import numpy as np
np.random.seed(1)
arr = np.random.randint(0, 100, 10)
print("Original array:", arr)
res = arr.argpartition(-3)
print("Indices:", res)
print("Three largest values:", arr[res[-3:]])
Результат выполнения кода:
"Original array: [37 12 72 9 75 5 79 64 16 1]"
"Indices: [9 5 1 8 3 0 6 7 2 4]"
"Three largest values: [79 64 72]"
Пример
Применим argpartition к двумерному массиву по заданной оси:
import numpy as np
arr = np.array([[3, 1, 4],
[1, 5, 9],
[2, 6, 5]])
print("Original array:")
print(arr)
res = arr.argpartition(1, axis=0)
print("Indices along axis 0:")
print(res)
Результат выполнения кода:
[[3 1 4]
[1 5 9]
[2 6 5]]
[[1 0 0]
[2 2 2]
[0 1 1]]
Пример
Используем метод для поиска позиций медианного элемента в массиве:
import numpy as np
arr = np.array([7, 2, 5, 1, 8, 4, 3, 6])
print("Original array:", arr)
median_idx = len(arr) // 2
res = arr.argpartition(median_idx)
print("Indices:", res)
print("Median value:", arr[res[median_idx]])
Результат выполнения кода:
"Original array: [7 2 5 1 8 4 3 6]"
"Indices: [3 1 6 5 2 7 0 4]"
"Median value: 4"