[LeetCode] 23. Merge k Sorted Lists (C++)

문제

배열 k가 정렬된 링크드리스트로 주어졌을 때, 각 링크드리스트를 하나의 오름차순으로 정렬된 링크드리스트로 나타내기

생각할 것

Min Heap을 사용해보기

각 리스트의 head를 모두 heap에 넣고, 제거하면서 연결하는 방법으로 풀어보았다.

💡 priority_queue

• 선언 방법 : priority_queue<자료형, container, 비교함수>
• 구조체를 priority queue에 넣는 방법

• struct cmp {
    bool operator()(const ListNode* n1, const ListNode* n2) {
        return n1 -> val > n2 -> val; // 오름차순 정렬 (일반적인 정렬기준과 반대)
    }
  };
  
  priority_queue<ListNode*, vector<ListNode*>, cmp> pq;

코드

/**
 * Definition for singly-linked list.
 * struct ListNode {
 *     int val;
 *     ListNode *next;
 *     ListNode() : val(0), next(nullptr) {}
 *     ListNode(int x) : val(x), next(nullptr) {}
 *     ListNode(int x, ListNode *next) : val(x), next(next) {}
 * };
 */
class Solution {
public:
    struct cmp {
        bool operator()(const ListNode* n1, const ListNode* n2) {
            return n1 -> val > n2 -> val; // 오름차순 정렬 (일반적인 정렬기준과 반대)
        }
    };

    ListNode* mergeKLists(vector<ListNode*>& lists) {
        priority_queue<ListNode*, vector<ListNode*>, cmp> pq;
        ListNode* resH = new ListNode();
        ListNode* res = resH;

        for(int i=0; i<lists.size(); i++){
            if(lists[i]) pq.push(lists[i]); // nullptr인 경우 힙에 넣지 않기 
        }

        while(!pq.empty()){
            res -> next = pq.top();
            pq.pop();
            res = res -> next; // 한 리스트의 노드가 연결되면 
            if(res -> next){ // 연결된 노드가 null이 아닐 시
                pq.push(res -> next); // min heap에 다음 노드를 넣기 
                // 주어진 리스트가 미리 정렬 되어있으므로 하나씩 넣어도 된다
            } 
        }

        return resH -> next; 
    }
};