代码练习4-合并 k 个升序的链表。

数据范围：节点总数 0≤𝑛≤5000，每个节点的val满足 ∣𝑣𝑎𝑙∣<=1000

要求：时间复杂度 O(nlogn)

输入：[{1,2,3},{4,5,6,7}]

输出：{1,2,3,4,5,6,7}

C++核心代码

#include <vector>
#include <queue>

struct ListNode {
    int val;
    ListNode *next;
    ListNode(int x) : val(x), next(nullptr) {}
};

class Solution {
public:
    struct compare {
        bool operator()(ListNode* a, ListNode* b) {
            return a->val > b->val;  // 优先队列中小的优先
        }
    };
    
    ListNode* mergeKLists(std::vector<ListNode*>& lists) {
        // 创建一个优先队列（最小堆）
        std::priority_queue<ListNode*, std::vector<ListNode*>, compare> pq;
        
        // 将每个链表的头节点加入优先队列
        for (auto list : lists) {
            if (list) {
                pq.push(list);
            }
        }
        
        // 创建一个虚拟头节点
        ListNode dummy(0);
        ListNode* current = &dummy;
        
        // 合并链表
        while (!pq.empty()) {
            // 获取最小节点并将其从堆中移除
            ListNode* node = pq.top();
            pq.pop();
            
            // 将最小节点添加到结果链表中
            current->next = node;
            current = current->next;
            
            // 将下一节点加入堆中
            if (node->next) {
                pq.push(node->next);
            }
        }
        
        // 返回合并后的链表的头节点
        return dummy.next;
    }
};

Python核心代码

import heapq
from typing import List

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

class Solution:
    def mergeKLists(self, lists: List[ListNode]) -> ListNode:
        # 定义一个最小堆
        heap = []
        
        # 将所有链表的头节点放入堆中
        for l in lists:
            if l:
                heapq.heappush(heap, (l.val, l))
        
        # 创建一个虚拟头节点
        dummy = ListNode(0)
        current = dummy
        
        # 合并链表
        while heap:
            val, node = heapq.heappop(heap)
            current.next = node
            current = current.next
            if node.next:
                heapq.heappush(heap, (node.next.val, node.next))
        
        return dummy.next