python - 有效地在 python 列表中查找 nlargest 元素的索引

Question

我们可以使用 heapq nlargest 来获取 nlargest 元素的值。我需要找到最大元素的索引。有什么建议么？

标签： python

Answer 1

import random
import heapq

values = list(range(10))
random.shuffle(values)
n = 3

pairs = heapq.nlargest(n, zip(values, range(len(values))))
indices = [i for value, i in pairs]

print(values)   # [0, 2, 8, 1, 5, 7, 4, 6, 3, 9]
print(pairs)    # [(9, 9), (8, 2), (7, 5)]
print(indices)  # [9, 2, 5]

这似乎只比使用enumerate. 这是一些计时代码：

setup = """
import random
import heapq
import timeit
import operator

values = list(range(10000))
random.shuffle(values)
n = 200

key = operator.itemgetter(1)
"""

stmt1 = """
pairs = heapq.nlargest(n, zip(values, range(len(values))))
indices = [i for value, i in pairs]
"""

stmt2 = """
pairs = heapq.nlargest(n, enumerate(values), key=key)
indices = [i for i, value in pairs]
"""

import timeit

print(timeit.timeit(stmt1, setup, number=1000))
print(timeit.timeit(stmt2, setup, number=1000))