如何使用 Python 堆实现多路归并问题？

import heapq

def merge_lists(sorted_lists):
    """
    合并多个有序列表
    """
    heap = []
    for i, lst in enumerate(sorted_lists):
        if lst:
            # 将列表的第一个元素及其所在列表索引作为元组加入堆中
            # 这里使用元组的原因是元组可以进行比较操作
            heapq.heappush(heap, (lst[0], i, 0))
    result = []
    while heap:
        # 弹出堆顶元素，该元素的第一个值是所在列表的元素值
        val, lst_idx, idx = heapq.heappop(heap)
        result.append(val)
        if idx+1 < len(sorted_lists[lst_idx]):
            # 将该元素所在列表的下一个元素加入堆中
            heapq.heappush(heap, (sorted_lists[lst_idx][idx+1], lst_idx, idx+1))
    return result

# 测试
lists = [[1, 3, 5], [2, 4, 6], [0, 7, 8]]
print(merge_lists(lists)) # [0, 1, 2, 3, 4, 5, 6, 7, 8]

str_lists = [['pidan', 'code', 'com'], ['pi', 'dan', 'code'], ['pi', 'dan', 'cheng', 'xia', 'com']]
print(merge_lists(str_lists)) # ['cheng', 'code', 'code', 'com', 'com', 'dan', 'dan', 'pi', 'pi', 'pidan', 'xia']