LC23. Merge k Sorted Lists

https://leetcode.com/problems/median-of-two-sorted-arrays/

  • 너무나 대표적인 문제.
  • 지난번에 실패했었는데, ListNode 자체를 넣으니까 진짜 쉽게 풀림
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
/**
 * Definition for singly-linked list.
 * struct ListNode {
 *     int val;
 *     ListNode *next;
 *     ListNode(int x) : val(x), next(NULL) {}
 * };
 */
 
class Solution {
public:
    ListNode* mergeKLists(vector<ListNode*>& lists) {
        auto comp = [](ListNode* a, ListNode* b) { return a->val > b->val; };
        priority_queue<ListNode*, vector<ListNode*>, decltype(comp)> pq(comp);
        for (auto& x: lists)
            if(x) pq.push(x);
        ListNode *head = nullptr, *tail = nullptr;
        while(pq.size()) {
            ListNode* cur = pq.top();
            pq.pop();
            if (cur->next) pq.push(cur->next);
            if(!head) {
                head = cur, tail = cur;
                continue;
            }
            tail->next = cur;
            tail = cur;
        }
        return head;
    }
};
  • 파이썬으로 Priority Queue 처음 써봄
  • 확실히 코드가 짧다.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
# Definition for singly-linked list.
# class ListNode:
#     def __init__(self, x):
#         self.val = x
#         self.next = None

from queue import PriorityQueue

class Node:
    def __init__(self, n):
        self.n = n
    def __lt__(self, other):
        return self.n.val < other.n.val

class Solution:
    def mergeKLists(self, lists: List[ListNode]) -> ListNode:
        pq = PriorityQueue()
        for h in lists:
            if h: pq.put(Node(h))
        head, tail = None, None
        while not pq.empty():
            cur = pq.get().n
            if cur.next: pq.put(Node(cur.next))
            if not head: head, tail = cur, cur
            else:
                tail.next = cur
                tail = cur
        return head
  • heapq 쓰는 버전, 다른 사람꺼를 참고해서 구현, 2배정도 빠름 위에꺼보다
  • 비교 불가능 예외가 발생하는것을 막기 위해서 i를 중간에 끼우는것은 나이스한듯
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
import heapq

class Solution:
    def mergeKLists(self, lists: List[ListNode]) -> ListNode:
        list_heap = [(n.val, i, n) for i, n in enumerate(lists) if n]
        heapq.heapify(list_heap)
        head, tail = None, None
        while len(list_heap):
            _, i, cur = heapq.heappop(list_heap)
            if cur and cur.next:
                heapq.heappush(list_heap, (cur.next.val, i, cur.next))
            if not head: head, tail = cur, cur
            else:
                tail.next = cur
                tail = cur
        return head