Поиск индекса N самых больших элементов в массиве Python/Список эффективно
Я сожалею заранее, если это дублированный вопрос, я искал эту информацию, но все еще не мог ее найти.
Можно ли организовать массив numpy (или список python), используя очень эффективные индексы N самых больших элементов в порядке убывания?
Например, массив:
a = array([4, 1, 0, 8, 5, 2])
Индексы крупнейших элементов в порядке убывания дадут (учитывая N = 6, все элементы включены):
8 → 3
5 → 4
4 → 0
2 → 5
1 → 1
0 → 2
result = [3, 4, 0, 5, 1, 2]
Я знаю, как сделать это, используя несколько глупый подход (например, сортировку массива и поиск каждого из N номеров для своих индексов), но мне было интересно, есть ли какая-либо эффективная библиотека, например, узкое место или heapq или, возможно, pythonic подход, чтобы сделать это очень быстро. Я должен применять его в нескольких массивах с элементами 300 тыс., Поэтому проблема в производительности.
Спасибо заранее!
UPDATE
Я прочитал ответы и решил использовать их с использованием 300k случайных целых чисел, вот результаты:
решение 1: sorted(range(len(a)), key=lambda i:a[i])
время: 230 мс
решение 2: heapq.nlargest(len(a), zip(a, itertools.count()))
время: 396 мс
решение 3: heapq.nlargest(len(a), enumerate(a), key=operator.itemgetter(1))
время: 864 мс
решение 4: def f(a,N): return np.argsort(a)[::-1][:N] (N = len(a))
время: 104 мс
Большое спасибо за быстрые и очень хорошие ответы!
Ответы
Ответ 1
Вы просмотрели встроенный метод numpy argsort
?:
http://docs.scipy.org/doc/numpy/reference/generated/numpy.argsort.html
Я могу сортировать массив с 300 000 случайных поплавков примерно на 29 мс на моей машине, используя этот метод.
def f(a,N):
return np.argsort(a)[::-1][:N]
Ответ 2
L = [4, 1, 0, 8, 5, 2]
sorted(range(len(L)), key=lambda i:L[i])
Ответ 3
Вы можете использовать heapq
, чтобы сделать это достаточно легко:
>>> heapq.nlargest(3, zip(a, itertools.count()))
[(8, 3), (5, 4), (4, 5)]
Кортежи сортируются путем сортировки по первому значению, затем второго и т.д. Это означает, что мы можем просто сделать кортеж (value, index)
и отсортировать, указав нам индексы значений (значения также указаны, но мы можем легко выбросить их).
Я использую zip()
и itertools.count()
, поскольку перечисление дает нам неправильный порядок, поэтому они будут отсортированы по индексу, а не по значению. В качестве альтернативы вы также можете сделать ((value, index) for index, value in enumerate(a))
, но я чувствую, что это менее понятно.
Другой вариант - дать ключ, сделав heapq.nlargest(3, enumerate(a), key=operator.itemgetter(1))
.
Ответ 4
Другой способ использования heapq
heapq.nlargest(n, range(len(a)), key=a.__getitem__)
Как прокомментировано в другом месте, он не будет бить сортировку, если он не очень большой и n<<len(a)
, потому что сортировка является относительно быстрой операцией в Python. Однако в конечном итоге медленный O (n) всегда будет бить O (n * log (n))