1. 程式人生 > 其它 >heapq模組中的兩個函式——nlargest()和nsmallest()

heapq模組中的兩個函式——nlargest()和nsmallest()

這兩個函式可以幫助我們在某個集合中找出最大或最小的N個元素。例如:

>>> import heapq
>>> nums=[1,8,2,23,7,-4,18,23,42,37,2]
>>> print(heapq.nlargest(3,nums))
[42, 37, 23]
>>> print(heapq.nsmallest(3,nums))
[-4, 1, 2]
如果正在尋找最大或者最小的N個元素,且同集合中元素的總數目相比,N很小,那麼下面這些函式就可以提供更好的效能。這些函式首先會在底層將資料轉化成列表,且元素會以堆得順序排列。例如:

>>> nums=[1,8,2,23,7,-4,18,23,42,37,2]
>>> import heapq
>>> heap=list(nums)
>>> heapq.heapify(heap)#將列表原地轉換成堆
>>> heap
[-4, 2, 1, 23, 7, 2, 18, 23, 42, 37, 8]
堆最重要的特性就是,heap[0]總是最小的那個元素,此外,接下來的最小元素可以一次通過heapq.heappop()的方法輕鬆的找到。該方法會將第一個元素(最小的)彈出,然後以第二小的元素取而代之。例如,要找到第三小的元素,可以這樣做:

>>> heapq.heappop(heap)
-4
>>> heapq.heappop(heap)
1
>>> heapq.heappop(heap)
2
另外,heapq.heappush()的功能是把一項值壓入堆,同時維持堆得排序要求。