LeetCode (C/C++, Java) 23. Merge k Sorted Lists

https://leetcode.com/problems/merge-k-sorted-lists/description/

 

---

풀이

현재 리스트의 맨 앞 값을 우선순위 큐에 저장해두고, 오름차순으로 정렬한다. 

큐에서 값을 하나씩 빼 root에 연결한다. 이 때, 다음 값보다 현재 값이 작을 때까지 값을 넣은 후, 리스트에 값이 남아 있으면 다시 큐에 넣어 순회한다.

 

더 빠른 풀이가 있길래 확인했는데, 리스트를 전부 순회하며 모든 수를 하나의 컨테이너에 넣고 정렬 후 이를 리스트 형태로 만들어 반환했다... 내부 소팅의 최적화가 잘 되어 있다는 걸 알려주고 싶은 문제였나 싶었다.

---

C/C++ 코드

class Solution {
public:
    typedef pair<int, ListNode*> ii;
    ListNode* mergeKLists(vector<ListNode*>& lists) {
        int k = lists.size();
        ListNode *root = NULL, *prev = NULL;
        priority_queue<ii, vector<ii>, greater<ii>> pq;

        if (!lists.size())
            return root;

        // 리스트의 첫 원소와 리스트 주소를 저장
        for (int i = 0; i < k; i++) {
            if (lists[i])
                pq.push({lists[i]->val, lists[i]});
        }

        while (!pq.empty()) {
            ListNode* now = pq.top().second;
            pq.pop();

            int nextMinVal = pq.top().first;

            // 현재보다 크고 전체에서 가장 작은 값 < 현재 노드의 값일 때까지 순회하며 리스트 생성
            while (now && now->val <= nextMinVal) {
                ListNode* nextNode = new ListNode(now->val);

                if (root == NULL)
                    root = nextNode;
                else
                    prev->next = nextNode;

                prev = nextNode;
                now = now -> next;
            }
            // 현재 리스트의 원소가 남아있다면 큐에 넣어 다시 탐색
            if (now)
                pq.push({now->val, now});
        }
        return root;
    }
};

 

Java 코드

class Solution {
    public ListNode mergeKLists(ListNode[] lists) {
        int k = lists.length;
        ListNode root = null, prev = null;
        PriorityQueue<ListNode> pq = new PriorityQueue<>((a, b) -> Integer.compare(a.val, b.val));

        if (k == 0) return root;

        for (int i = 0; i < k; i++) {
            if (lists[i] != null) {
                pq.add(lists[i]);
            }
        }

       while (!pq.isEmpty()) {
            ListNode now = pq.poll(); // 최소값을 가진 노드를 꺼냄

            if (root == null)
                root = now;
            else
                prev.next = now;

            prev = now;

            if (now.next != null) 
                pq.add(now.next);
        }
        return root;
    }
}