合并K个有序链表

news2024/9/9 4:58:19

题目

给你一个链表数组，每个链表都已经按升序排列。

请你将所有链表合并到一个升序链表中，返回合并后的链表。

示例1：

输入： $lists=[[1,4,5],[1,3,4],[2,6]]$

输出： $[1,1,2,3,4,4,5,6]$
示例2：

输入： $lists = []$

输出： $[]$

示例3：

输入： $lists=[[]]$

输出： $[]$

解题

优先队列（最小堆）法

优先队列可以帮助我们每次从所有链表的头节点中找到最小的节点，将其放入合并后的链表中，然后从相应的链表中移除这个节点。具体步骤如下：

初始化优先队列：

对每个链表的头节点进行初始化，将其放入优先队列中。优先队列按节点的值进行排序。

构建合并后的链表：

每次从优先队列中取出最小节点，将其加入结果链表。
如果该节点有下一个节点，则将下一个节点放入优先队列中。

返回结果：

最终，返回合并后的链表头节点。

from heapq import heappop, heappush


class ListNode:
    def __init__(self, val=0, next=None):
        self.val = val
        self.next = next


def mergeKLists(lists):
    # 创建一个虚拟头节点
    dummy = ListNode(0)
    current = dummy

    # 定义一个最小堆
    heap = []

    # 初始化堆，将每个链表的头节点放入堆中
    for i, l in enumerate(lists):
        if l:
            heappush(heap, (l.val, i, l))

    # 合并链表
    while heap:
        val, i, node = heappop(heap)
        current.next = node
        current = current.next
        if node.next:
            heappush(heap, (node.next.val, i, node.next))

    return dummy.next


# 辅助函数：创建链表
def create_linked_list(values):
    if not values:
        return None
    head = ListNode(values[0])
    current = head
    for value in values[1:]:
        current.next = ListNode(value)
        current = current.next
    return head


# 辅助函数：打印链表
def print_linked_list(head):
    values = []
    while head:
        values.append(head.val)
        head = head.next
    print(" -> ".join(map(str, values)))


# 示例测试
if __name__ == '__main__':
    lists = [
        create_linked_list([1, 4, 5]),
        create_linked_list([1, 3, 4]),
        create_linked_list([2, 6])
    ]
    merged_head = mergeKLists(lists)
    print("合并后的链表:")
    print_linked_list(merged_head)