查找列表中最高的n个元素及其位置。蟒蛇

时间:2013-02-28 13:23:22

标签: python performance list

最近我被要求编写一些代码来查找列表中最高的n个元素,并返回值和位置。

你能比这更快(在执行时间方面)吗?

def highest(L, n):
    return sorted(enumerate(L), reverse=True, key=lambda x: x[1])[:n]

if __name__ == '__main__':

    M = [102, 56, 2355, 3, 25, 78, 19, 25, 1002, -54, 0, 23, -1]
    r = highest(M,5)
    print r  #[(2, 2355), (8, 1002), (0, 102), (5, 78), (1, 56)]

2 个答案:

答案 0 :(得分:9)

如果n与列表的长度相比较小,heapq.nlargest应该比排序整个列表更快。它也更具可读性。

def highest(L, n):
    return heapq.nlargest(n, enumerate(L), key=operator.itemgetter(1))

>>> M = [102, 56, 2355, 3, 25, 78, 19, 25, 1002, -54, 0, 23, -1]
>>> highest(M,5)
[(2, 2355), (8, 1002), (0, 102), (5, 78), (1, 56)]

这将在O(N + nlogn)中工作,其中N是列表的长度,n是要返回的项目数,而不是用于排序的O(NlogN)。 / p>

答案 1 :(得分:1)

只需添加一个,您可以在pandas中使用kth_smallest来查找O(N)中的值。

import numpy as np
import pandas as pd
a = np.array([102, 56, 2355, 3, 25, 78, 19, 25, 1002, -54, 0, 23, -1.0])
pd.algos.kth_smallest(a, len(a)-5)

代码在这里:

https://github.com/pydata/pandas/blob/master/pandas/algos.pyx#L653

注意:kth_smallest仅返回值,但您可以扫描数组以查找位置。