使用priority queue合并K链表时比较器未对所有元素生效问题排查

使用优先队列合并K个链表时，发现比较器并未对所有元素生效。当所有链表仅包含一个节点时代码可以正常运行，但多节点链表的场景下输出结果不符合预期。

代码示例 #

#include<iostream>
#include<vector>
#include<queue>
using namespace std;
class ListNode {
    public:
        int val;
        ListNode *next;
        ListNode() {
            this->val = 0;
            this->next = nullptr;
        }
        ListNode(int x) {
            this->val = x;
            this->next = nullptr;
        }
        
        ListNode(int x, ListNode *next):val(x),next(next) {
        }
};

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> pq;
        ListNode *dummy = new ListNode(0);
        ListNode *cur = dummy;
        ListNode *top;
        for (int i=0; i< lists.size(); i++) {
        pq.push(lists[i]);
    }
    while(!pq.empty()) {
            top = pq.top();
        pq.pop();
            cur->next = top;
            cur = cur->next;
        cur->next = nullptr;
            if(top->next != nullptr) {
                pq.push(top->next);
            }
        }
        return dummy->next;
    }
};

int main() {
    Solution obj;
    ListNode *a = new ListNode(9);
    a->next = new ListNode(34);
    a->next->next = new ListNode(55);   
    ListNode *b = new ListNode(45);
    b->next = new ListNode(4);
    b->next->next = new ListNode(5);
    vector<ListNode*> Lists;
    Lists.push_back(a);
    Lists.push_back(b);
    ListNode *result = obj.mergeKLists(Lists);
    while(result != nullptr) {
        cout<<result->val<< " ";
        result = result->next;
    }
    cout<<endl;
    return 0;
}

输出结果 #

9 34 45 4 5 55

问题分析 #

你的代码逻辑本身没有问题，问题出在输入链表不符合算法的前提假设。

优先队列合并K个链表的经典解法，要求每个输入链表本身是有序的（通常为升序）。你的第二个链表b是45 -> 4 -> 5，本身是无序的：45之后跟着更小的4，导致这个4只能在45被取出并处理后，才会被加入优先队列，自然会出现在45之后，最终导致整体结果无序。

当所有链表仅含一个节点时，输入天然是“有序”的（单个节点没有顺序问题），所以代码能正常运行；但多节点链表无序时，就会出现你看到的现象。

解决方案 #

1. 如果需求是合并无序的K个链表：优先队列方法不适用，需要先把所有链表的节点收集到容器中，排序后再重新构建链表。
2. 如果需求是合并有序的K个链表：修正输入链表的顺序即可。比如把第二个链表调整为升序4 -> 5 -> 45，此时代码会输出正确的升序结果：4 5 9 34 45 55。

另外补充一个细节：代码未处理输入链表为空的情况（比如lists中存在nullptr），如果有这种场景，需在入队前添加判断：

for (int i=0; i< lists.size(); i++) {
    if(lists[i] != nullptr) { // 跳过空链表
        pq.push(lists[i]);
    }
}

内容的提问来源于stack exchange，提问作者BEPP