Python中的堆和優先隊列

分別是heapq和queue.PriorityQueue這兩個模塊

import heapq        

from queue import PriorityQueue as PQ
PriorityQueue模塊定義如下所示:
class PriorityQueue(Queue):
    '''Variant of Queue that retrieves open entries in priority order (lowest first).

    Entries are typically tuples of the form:  (priority number, data).
    '''

    def _init(self, maxsize):
        self.queue = []

    def _qsize(self):
        return len(self.queue)

    def _put(self, item):
        heappush(self.queue, item)

    def _get(self):
        return heappop(self.queue)
他含有一個屬性queue,輸出隊列中每個元素,三個方法,分別是qsize(),代表優先級隊列中元素個數,put(),用heappush()方法將元素放入隊列,get(),用heappop()方法將元素從隊列取出。注意這個庫在解答leetcode相關問題時(比如top-K問題,leetcode-215,347等)不能用。
例子如下
from queue import PriorityQueue as PQ
pq = PQ()
pq.put((-2, 'c'))
pq.put((-3, 'b'))
pq.put((-4,'a'))
pq.queue # output: [(-4, 'a'), (-2, 'c'), (-3, 'b')]
#只能是最小堆,所以按照元素從小到大(如果元素是list或者tuple,那麼以第一個元素的順序)的順序出隊
pq.get() # output: (-4, 'a')
#每當堆首元素出隊後,剩下的元素又會重新排成堆
pq.get() #output: [(-3, 'b'), (-2, 'c')]
pq.get() # output: (-3,'b')
pq.get() # output: (-2,'c')
#最後結果爲空
pq.queue # output:[]
heappush模塊代碼如下,功能:將item添加到當前最小堆最後,並再恢復成最小堆(只是把元素插到堆尾再調整,之前已經是堆了,現在只有最後一個元素可能不滿足堆得要求,調整也只會沿着一條路徑自底向上,時間複雜度O(lgn))
def heappush(heap, item):
    """Push item onto heap, maintaining the heap invariant."""
    heap.append(item)
    _siftdown(heap, 0, len(heap)-1)


def _siftdown(heap, startpos, pos):
newitem = heap[pos]
# Follow the path to the root, moving parents down until finding a place
# newitem fits.
while pos > startpos:
    parentpos = (pos - 1) >> 1
    parent = heap[parentpos]
    if newitem < parent:
        heap[pos] = parent
        pos = parentpos
        continue
    break
heap[pos] = newitem
heappop模塊代碼如下,功能:當前最小堆中首元素,再將當前堆尾元素放在堆首,之後將這個列表再回覆成堆。(只是把元素插到堆首再調整,之前已經是堆了,現在只有第一個元素可能不滿足堆得要求,調整也只會沿着一條路徑從根節點向下,時間複雜度O(lgn))
def heappop(heap):
    """Pop the smallest item off the heap, maintaining the heap invariant."""
    lastelt = heap.pop()    # raises appropriate IndexError if heap is empty
    if heap:
        returnitem = heap[0]
        heap[0] = lastelt
        _siftup(heap, 0)
        return returnitem
    return lastelt

def _siftup(heap, pos):
    endpos = len(heap)
    startpos = pos
    newitem = heap[pos]
    # Bubble up the smaller child until hitting a leaf.
    childpos = 2*pos + 1    # leftmost child position
    while childpos < endpos:
        # Set childpos to index of smaller child.
        rightpos = childpos + 1
        if rightpos < endpos and not heap[childpos] < heap[rightpos]:
            childpos = rightpos
        # Move the smaller child up.
        heap[pos] = heap[childpos]
        pos = childpos
        childpos = 2*pos + 1
    # The leaf at pos is empty now.  Put newitem there, and bubble it up
    # to its final resting place (by sifting its parents down).
    heap[pos] = newitem
    _siftdown(heap, startpos, pos)

def _siftdown(heap, startpos, pos):
    newitem = heap[pos]
    # Follow the path to the root, moving parents down until finding a place
    # newitem fits.
    while pos > startpos:
        parentpos = (pos - 1) >> 1
        parent = heap[parentpos]
        if newitem < parent:
            heap[pos] = parent
            pos = parentpos
            continue
        break
    heap[pos] = newitem
heapfiy模塊代碼如下,功能:在O(len(x))時間複雜度將列表x原地轉換成最小堆
def heapify(x):
    """Transform list into a heap, in-place, in O(len(x)) time."""
    n = len(x)
    # Transform bottom-up.  The largest index there's any point to looking at
    # is the largest with a child index in-range, so must have 2*i + 1 < n,
    # or i < (n-1)/2.  If n is even = 2*j, this is (2*j-1)/2 = j-1/2 so
    # j-1 is the largest, which is n//2 - 1.  If n is odd = 2*j+1, this is
    # (2*j+1-1)/2 = j so j-1 is the largest, and that's again n//2-1.
    for i in reversed(range(n//2)):
        _siftup(x, i)
這個某塊可以直接使用使用,通過import heapq,所以當在leetcode中遇到top-K或類似問題需要用到優先隊列時(Python),不需要用優先隊列了,直接用堆即可,不過注意python中的heapq只是最小堆,如果需要最大堆,那麼在把元素入堆的時候需要加負號。並且最後輸出結果的時候也要加負號,然後在逆序輸出。

下面以leetcode-215和leetcode-347題目的用優先隊列和堆的解法來進一步瞭解兩者的關係。

leetcode-215

leetcode-215

class Solution:
    def findKthLargest(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: int
        """
        import heapq
        #heapq中直接有這個函數
        return heapq.nlargest(k, nums)[-1]


# 用堆,時間複雜度O(N + klog(N))
class Solution:
    def findKthLargest(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: int
        """
        import heapq
        # heapify只能創建一個最小堆,所以加負號,但根據題目意思應該用最大堆
        nums = [-num for num in nums]
        heapq.heapify(nums)
        res = float('inf')
        for _ in range(k):
            res = heapq.heappop(nums)
        return -res
leetcode-347

在這裏插入圖片描述

# top-K問題常用解題思路,用優先隊列,時間複雜度O(nlgk)
# 維護一個含有k個元素的優先隊列。如果遍歷到的元素比隊列中的最小頻率元素的頻率高,則取出隊列中最小頻率
# 的元素,將新元素入隊。最終,隊列中剩下的,就是前k個出現頻率最高的元素。
class Solution(object):
    def topKFrequent(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: List[int]
        """
        freq = {}
        for num in nums:
            if num not in freq:
                freq[num] = 1
            else:
                freq[num] += 1
        # 掃描數組,維護當前出現頻率最高的k個元素
        # 在優先隊列中,按照頻率排序,所以數據對是(頻率,元素)的形式
        # python的堆只有最小堆,優先隊列就是用堆實現的,所以優先隊列也只有最小優先隊列
        from queue import PriorityQueue as PQ
        pq = PQ()
        for v, f in freq.items():
            if pq.qsize() == k:
                (fr, va) = pq.get()
                # pq只是最小優先隊列,爲了實現最大優先隊列,所以需要取負數
                if -f > -fr:
                    pq.put((-f, v))
                else:
                    pq.put((fr, va))
            else:
                pq.put((-f, v))
        res = []
        while pq.queue:
            res.append((pq.get()[1]))
        return res
    
    
 #上面用優先隊列只是爲了瞭解這個模塊中優先隊列的用法,實際上用起來也不方便,可以看到上面寫的很繁瑣,
 不易懂,而且在leetcode上面提交還會顯示ImportError: No module named queue


class Solution(object):
    def topKFrequent(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: List[int]
        """
        # 上面的優先隊列就是用heapq實現的,所以這裏直接使用heapq堆
        import heapq
        freq = {}
        res = []
        for num in nums:
            if num not in freq:
                freq[num] = 1
            else:
                freq[num] += 1
        # 1.我們需要按照數字出現頻率進行排序,所以val在前,key在後
        # 2.我們需要最大堆,但是heapq只實現了最小堆,所以加個負號,模擬最大堆
        max_heap = [(-val, key) for key, val in freq.items()]
        heapq.heapify(max_heap)
        for i in range(k):
            res.append(heapq.heappop(max_heap)[1])
        return res
    
    
    nums = [1, 1, 1, 1, 2, 2, 2, 3, 3, 4, 4, 5, 6, 7]
    k = 4
    print(Solution().topKFrequent(nums, k))

參考

最小堆 構建、插入、刪除的過程圖解
排序算法總結-堆排序

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章