[LeetCode, 리트코드] Merge k Sorted Lists

1. Problem

정렬된 k개의 링크드 리스트를 하나의 정렬된 리스트로 합쳐라.

Example:

Input:
[
  1->4->5,
  1->3->4,
  2->6
]
Output: 1->1->2->3->4->4->5->6

2. Solution

처음에는 각 리스트들의 작은값을 찾아서 하나하나 비교해보는 풀이 방법을 택했다. 그러나 다른 솔루션을 보니 heap을 이용하여 이 부분을 최적화시킬 수 있는 아이디어가 있어서 heap을 이용하여 다시 코드를 작성했다. 

heap을 이용하여 최적화시키는 부분은 이렇다. k개의 리스트 중 맨 앞쪽의 가장 작은 요소를 구하기 위해서 기존의 heap에 있던 최상위 노드와 비교를 한다. 이 최상위 노드는 기존에 비교했었던 요소들 중 가장 작은 요소다. 

이것을 이용하여 O(kn)이었던 시간복잡도를 O(lgk*N)으로 줄일 수 있다.

3. Code1 - 하나씩 비교하기

class Solution {
public:
    ListNode* mergeKLists(vector<ListNode*>& lists) {
        ListNode* ret = new ListNode(-9999);
        ListNode* curr = ret;
        
        merge(lists, curr);
        return ret->next;
    }
    
    void merge(vector<ListNode*>& lists, ListNode* curr){
        ListNode* min = new ListNode(987654321);
        int min_idx;
        for(int i=0; i < lists.size(); ++i){
            if(lists[i] == nullptr)
                continue;
            if(lists[i]->val <= min->val){
                min = lists[i];
                min_idx = i;
            }
        }

        if(min->val != 987654321){
            curr->next = min;
            lists[min_idx] = lists[min_idx]->next;
            curr = curr->next;
            merge(lists, curr);
        }
        return;
    }
};

4. Code2 - Heap을 이용한 최적화

class Solution {
public:
    struct comparator {
        bool operator()(ListNode* a, ListNode* b){
            return a->val > b->val;
        }
    };
    
    ListNode* mergeKLists(vector<ListNode*>& lists) {
        priority_queue<ListNode*, vector<ListNode*>, comparator> pq;
        for(auto node : lists)
            if(node)
                pq.push(node);
        ListNode* ret = new ListNode(-99999);
        ListNode* cur = ret;
        
        while(pq.empty() == false){
            cur->next = pq.top();
            pq.pop();
            if(cur->next->next) pq.push(cur->next->next);
            cur = cur->next;
        }
        
        return ret->next ? ret->next : nullptr;
    }
};