有效地在python列表中找到最大元素的索引

时间:2019-04-15 19:20:30

标签: python

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

1 个答案:

答案 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))