python【heapq】&& leetcode 23. Merge k Sorted Lists

heapq模塊

heapq 模塊是python裏用來實現 ——最小堆 ，又被稱爲優先隊列算法，官方文檔。

最近用python刷leetcode用的比較多，用一些例子做個筆記。

創建堆 - 最小堆

單個添加創建堆 - heappush

import heapq

data = [1,5,3,2,8,5]
heap = []

for n in data:
    heapq.heappush(heap, n)

對已存在的序列轉化爲堆 - heapify

import heapq

data = [1,5,3,2,8,5]
heapq.heapify(data)

對多個序列轉化爲堆 - merge

import heapq

num1 = [32, 3, 5, 34, 54, 23, 132]
num2 = [23, 2, 12, 656, 324, 23, 54]
num1 = sorted(num1)
num2 = sorted(num2)

res = heapq.merge(num1, num2)
print(list(res))  # [2, 3, 5, 12, 23, 23, 23, 32, 34, 54, 54, 132, 324, 656]

創建堆 - 最大堆

由於heapq默認爲最小堆，我們可以通過將數組元素取相反數， 然後在取出堆頂時進行取反, 即可獲得原值。

import heapq

data = [1, 5, 3, 2, 8, 5]
li = []
for i in data:
    heapq.heappush(li, -i)

print(li)  # [-8, -5, -5, -1, -2, -3]
print(-li[0])  # 8`

訪問堆內容

查看最小值 - [0]

import heapq
data = [1, 5, 3, 2, 8, 5]
heapq.heapify(data)
print(data[0]) # 1

彈出最小值 - heappop

會改變原數據, 類似於列表的 pop

import heapq

data = [1, 5, 3, 2, 8, 5]
heapq.heapify(data)

print(heapq.heappop(data))  # 1
print(data)  # [2, 5, 3, 5, 8]

向堆內推送值 - heappush

和上面創建堆例子類似。

彈出最小值並加入一個值 - heappushpop

import heapq

data = [1, 5, 3, 2, 8, 5]
heapq.heapify(data)

print(heapq.heappushpop(data, 1))  # 1
print(data)  # [1, 2, 3, 5, 8, 5]

彈出最小值並加入一個值 - heapreplace

彈出最小值, 添加新值, 且自動排序保持是最小堆

是 heappushpop 的高效版本, 在py3 中適用。

import heapq

data = [1, 5, 3, 2, 8, 5]
heapq.heapify(data)

print(heapq.heapreplace(data, 10))  # 1
print(data)  # [2, 5, 3, 10, 8, 5]

k 值問題 - nlargest / nsmallest

找出堆中最小 / 大的 k 個值

import heapq

data = [1, 5, 3, 2, 8, 5]

heapq.heapify(data)
print(data)  # [1, 2, 3, 5, 8, 5]
print(heapq.nlargest(2, data))  # [8, 5]
print(heapq.nsmallest(2, data))  # [1, 2]

可以接收一個參數 key , 用於指定選項進行選取

用法類似於 sorted 的 key

import heapq
from pprint import pprint
portfolio = [
    {'name': 'IBM', 'shares': 100, 'price': 91.1},
    {'name': 'AAPL', 'shares': 50, 'price': 543.22},
    {'name': 'FB', 'shares': 200, 'price': 21.09},
    {'name': 'HPQ', 'shares': 35, 'price': 31.75},
    {'name': 'YHOO', 'shares': 45, 'price': 16.35},
    {'name': 'ACME', 'shares': 75, 'price': 115.65}
]
cheap = heapq.nsmallest(3, portfolio, key=lambda s: s['price'])
expensive = heapq.nlargest(3, portfolio, key=lambda s: s['price'])
pprint(cheap)
pprint(expensive)

"""
輸出：
[{'name': 'YHOO', 'price': 16.35, 'shares': 45},
 {'name': 'FB', 'price': 21.09, 'shares': 200},
 {'name': 'HPQ', 'price': 31.75, 'shares': 35}]
[{'name': 'AAPL', 'price': 543.22, 'shares': 50},
 {'name': 'ACME', 'price': 115.65, 'shares': 75},
 {'name': 'IBM', 'price': 91.1, 'shares': 100}]
"""

優先隊列：

heapq堆的鍵值可以是元組，比較的機制是從元組首位[0]開始,可以實現帶權值的數據類型進行排序。

>>> h = []
>>> heappush(h, (5, 'write code'))
>>> heappush(h, (7, 'release product'))
>>> heappush(h, (1, 'write spec'))
>>> heappush(h, (3, 'create tests'))
>>> heappop(h)
(1, 'write spec')

在實現優先隊列的時候，難免有兩個數據權值重複，我們這時候可以添加一個位置表示入隊列順序，增加一個可以比較的優先級。

class PriorityQueue:
    def __init__(self):
        self.__queue = []
        self.__index = 0
        
    def push(self, item, priority):
        heapq.heappush(self.__queue, (-priority, self.__index, item))
        # 第一個參數：添加進的目標序列
        # 第二個參數：將一個元組作爲整體添加進序列，目的是爲了方便比較
        # 在priority相等的情況下，比較_index
        # priority爲負數使得添加時按照優先級從大到小排序，因爲堆排序的序列的第一個元素永遠是最小的
        self.__index += 1
        
    def pop(self):
        # 返回按照-priority 和 _index 排序後的第一個元素(是一個元組)的最後一個元素(item)
        return heapq.heappop(self.__queue)[-1]
 
q = PriorityQueue()
q.push("bar", 2)
q.push("foo", 1)
q.push("gork", 3)
q.push("new", 1)
print(q.pop())
print(q.pop())
print(q.pop())
print(q.pop())
 
"""
gork  # 優先級最高
bar   # 優先級第二
foo   # 優先級與new相同，比較index，因爲先加入，index比new小，所以排在前面
new
"""

Leetcode 23. Merge k Sorted Lists

合併K個有序的鏈表，總共N個數。

思路

最暴力的做法是把N個數放到一個數組裏，再一起排序，O(NlogN)。

當然由於k個鏈表是有序的，我們實際上只需要維護k個指針從k個鏈表的頭向尾滑動，每次選取k個鏈表的表頭裏的最小加入新的有序鏈表裏。

這裏我們就可以借用最小堆(優先隊列)維護k個鏈表的當前頭位置的值。

時間複雜度就變成O(N*logK)

這裏我利用python自帶heapq實現最小堆的時候，直接把（node.value,node)組成元組放進隊列裏，結果leetcode編譯器報錯,顯示<無法對LISTNode類型數據進行操作…服了…

class Solution(object):
    def mergeKLists(self, lists):
        import heapq
        que = []
        for node in (lists):
            if node != None : 
                heapq.heappush(que ,(node.val, node))
        dummy_node = ListNode(-1)
        cur = dummy_node
        while que:
            val, node =  heapq.heappop(que)
            cur.next = node
            cur = cur.next
            if node.next != None:
                heapq.heappush(que, (node.next.val, node.next))
        return dummy_node.next

但更坑的是我發現雖然python3提交編譯器報錯，但python提交可以通過…

好在python3裏heapq可以對編譯器自有變量類型組成的元組進行比較。

同時查資料發現元組在heapq裏比較的機制是從元組首位0開始，即遇到相同，就比較元組下一位，比如(1,2), (1,3)，前者比後者小。

而這題剛好node值有重複的，同時ListNode無法被比較，所以編譯器報錯。

於是把代碼修改一下，存的是節點值和節點的索引，都是int類型，當node.val相同時，索引小的先彈出。

class Solution:
    def mergeKLists(self, lists: List[ListNode]) -> ListNode:
        import heapq
        que , curs = [] , [] # curs存K個鏈表滑動的頭指針
        for index, node in enumerate(lists):
                if node!=None:
                    heapq.heappush(que ,(node.val, index))
                curs.append(node)

        dummy_node = ListNode(-1)
        cur = dummy_node
        while que:
            val, index =  heapq.heappop(que)
            cur.next = lists[index]
            cur = cur.next
            lists[index] = lists[index].next
            if lists[index] != None:
                heapq.heappush(que, (lists[index].val, index))
        return dummy_node.next

還是C++強大，任何數據類型都可以對<重載一下，都可以實現排序，即我們只要重寫一個cmp函數，丟到優先隊列即可。
C++

class Solution {
public:
    struct cmp{  //對新的數據類型的<進行重寫
       bool operator()(ListNode *a,ListNode *b){
          return a->val > b->val;
       }
    };
    ListNode* mergeKLists(vector<ListNode*>& lists) {
       priority_queue<ListNode* ,vector<ListNode*> , cmp> heapk;
       for(auto p:lists){
        if(p!=NULL){
            heapk.push(p);
        }
       }
       ListNode *pHead = new ListNode(-1);
       ListNode *pCur = pHead;
       while(!heapk.empty()){
          ListNode *top = heapk.top();heapk.pop();
          pCur->next = top;
          pCur = pCur->next;
          if(top->next!=NULL){
            heapk.push(top->next);
          } 
       }
       pCur = pHead->next;
       delete pHead;
       return pCur;
    }
};