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& lists) {
priority_queue , 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;
}
};