我们可以使用heapq nlargest来获取nlargest元素的值。 我需要找到最大元素的索引。 有什么建议吗?
答案 0 :(得分: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))