在Python中获取列表中较小的n个元素

Question

我需要在Python中获得较少的n个列表.我需要这个非常快,因为它是性能的关键部分,需要重复很多次.

n通常不大于10,列表通常有大约20000个元素.每次调用该函数时,列表总是不同的.无法进行排序.

最初,我写了这个函数:

def mins(items, n):
    mins = [float('inf')]*n
    for item in items:
        for i, min in enumerate(mins):
            if item < min:
                mins.insert(i, item)
                mins.pop()
                break
    return mins

但是这个函数无法击败对整个列表进行排序的简单排序(项目)[:n].这是我的测试:

from random import randint, random
import time

test_data = [randint(10, 50) + random() for i in range(20000)]

init = time.time()
mins = mins(test_data, 8)
print 'mins(items, n):', time.time() - init

init = time.time()
mins = sorted(test_data)[:8]
print 'sorted(items)[:n]:', time.time() - init

结果:

mins(items, n): 0.0632939338684
sorted(items)[:n]: 0.0231449604034

sorted()[:n]快三倍.我相信这是因为:

1. insert()操作成本很高,因为Python列表不是链表.
2. sorted()是一个优化的c函数,我的是纯python.

有没有办法击败sorted()[:n]？我应该使用C扩展,Pyrex或Psyco或类似的东西吗？

提前感谢您的回答.

Answer 1

你实际上想要一个排序的分钟序列.

mins = items[:n]
mins.sort()
for i in items[n:]:
    if i < mins[-1]: 
        mins.append(i)
        mins.sort()
        mins= mins[:n]

这样运行速度要快得多,因为你甚至没有看到分钟,除非它可以证明它的值大于给定的项目.大约是原算法时间的十分之一.

我的戴尔零时间运行.我必须运行10次以获得可测量的运行时间.

mins(items, n): 0.297000169754
sorted(items)[:n]: 0.109999895096
mins2(items)[:n]: 0.0309998989105

使用bisect.insort而不是追加和排序可以进一步提高头发.

Answer 2

import heapq

nlesser_items = heapq.nsmallest(n, items)

这是S.Lott算法的正确版本:

from bisect    import insort
from itertools import islice

def nsmallest_slott_bisect(n, iterable, insort=insort):
    it   = iter(iterable)
    mins = sorted(islice(it, n))
    for el in it:
        if el <= mins[-1]: #NOTE: equal sign is to preserve duplicates
            insort(mins, el)
            mins.pop()

    return mins

性能:

$ python -mtimeit -s "import marshal; from nsmallest import nsmallest$label as nsmallest; items = marshal.load(open('items.marshal','rb')); n = 10"\
 "nsmallest(n, items)"

nsmallest_heapq
100 loops, best of 3: 12.9 msec per loop
nsmallest_slott_list
100 loops, best of 3: 4.37 msec per loop
nsmallest_slott_bisect
100 loops, best of 3: 3.95 msec per loop

nsmallest_slott_bisect是快3倍比heapq的nsmallest(对于n = 10,LEN(项目)= 20000).nsmallest_slott_list只是略微慢一点.目前还不清楚为什么heapq的最小是如此缓慢; 它的算法几乎与上面给出的相同(对于小n).