Блог

У меня есть массив [2, 4, 2, 1, 3, 2, 2] , и я хочу получить верхние 3 индекса значений в порядке убывания. Таким образом, мой вывод должен быть [1,4,0], которые являются индексами верхних 3 элементов массива в порядке убывания. (Допускаются аналогичные индексы элементов). Я использую bottelneck, который частично выполняет мою работу.

 import bottleneck
import numpy as np
a = np.array([2, 4, 2, 1, 3, 2, 2])
b = np.array(bottleneck.argpartsort(-a, 3)[:3])
print(b)
  

что дает [4 1 5] правильные верхние 3 (n) индекса массива, но не в порядке убывания. Есть ли какая-либо встроенная функция для настройки этого вывода для обеспечения порядка убывания

Вы можете использовать np.argsort

 import numpy as np
a = np.array([2, 4, 2, 1, 3, 2, 2])
b = np.argsort(a)[-3:]
b = b[::-1]
  

Примечание: Последние две строки могут быть дополнительно упрощены до (хотя и с небольшим вычислительным преимуществом)

 b = np.argsort(a)[-3:][::-1]
# or, as juanpa suggested
b = np.argsort(a)[-1:-4:-1]
  

Все они дают следующее значение b

 array([1, 4, 6])
  

Я полагаю, что bottleneck.argpartsort делает что-то похожее на np.argpartiton . Итак, давайте воспользуемся этим. Теперь, np.argpartiton вместо 3 , нам нужно использовать range(3) для сохранения отсортированного порядка. Используя только скаляр, argpartiton не будет сортировать эти 3 элементы, но с range(3) это будет.

Давайте приведем пример, чтобы продемонстрировать это —

 In [360]: a
Out[360]: 
array([ 0.6082239 ,  0.74932587,  0.50574681,  0.85241966,  0.91742228,
        0.9984438 ,  0.6876266 ,  0.90651069,  0.53746683,  0.70674607])

In [361]: np.argpartition(-a,3) # 4,5,7 are the top 3, but not in sorted order
Out[361]: array([4, 5, 7, 3, 1, 9, 6, 2, 8, 0])

In [362]: np.argpartition(-a,range(3)) # Sorted order maintained for top 3
Out[362]: array([5, 4, 7, 3, 1, 0, 6, 2, 8, 9])
  

Кроме того, с np.argpartiton это не приведет к сортировке оставшихся элементов в массиве и, таким образом, обеспечит преимущество во время выполнения. Это было бы чрезвычайно полезно, если мы хотим получить верхние N элементы в отсортированном порядке, где N это относительно меньшее число, чем длина массива.

Таким образом, у нас была бы реализация с np.argpartiton , примерно такая —

 a[np.argpartition(-a,range(3))[:3]]
  

Тест во время выполнения —

 In [342]: a = np.random.rand(10000)

In [343]: %timeit a[np.argsort(a)[-1:-4:-1]] #@R. S. Nikhil Krishna's soln
1000 loops, best of 3: 907 µs per loop

In [344]: %timeit a[np.argpartition(-a,range(3))[:3]]
10000 loops, best of 3: 67.9 µs per loop