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;
    }
};