LeetCode) mergeKLists

 

내가 푼 코드 제공받은 리스트 중에 가장 작은 값을 비교하여 새로운 노드에 값을 계속 추가해줬다.

근데 런타임이 아주 길게 나오며 아슬아슬하게 통과한 것 같다.

class ListNode:
    def __init__(self, val=0, next=None):
        self.val = val
        self.next = next
from typing import List, Optional

class Solution:
    def mergeKLists(self, lists: List[Optional[ListNode]]) -> Optional[ListNode]:
        for i in range(len(lists)-1,-1,-1):
            if not lists[i]:
                del lists[i]

        if len(lists) == 0:
            return

        def find_min_val(lists):
            min_i = 10001
            min_i_index = 0
            for i in range(len(lists)):
                if lists[i]:
                    if min_i > lists[i].val:
                        min_i = lists[i].val
                        min_i_index = i

            return min_i_index

        min_index = find_min_val(lists)
        result = lists[min_index]

        if lists[min_index].next:
            lists[min_index] = lists[min_index].next
        else:
            del lists[min_index]
        node = result
        while lists:
            min_index = find_min_val(lists)
            node.next = lists[min_index]

            if lists[min_index].next:
                lists[min_index] = lists[min_index].next
            else:
                del lists[min_index]
            node = node.next

        return result

 

 

답안 코드는 heapq를 사용하여 효율적으로 문제를 해결했다. 자료구조 heap을 사용하는 것을 항상 생각해야겠다.

import heapq
class ListNode:
    def __init__(self, val=0, next=None):
        self.val = val
        self.next = next
from typing import List, Optional

class Solution:
    def mergeKLists(self, lists: List[Optional[ListNode]]) -> Optional[ListNode]:
        root = result = ListNode(None)
        heap = []

        for i in range(len(lists)):
            if lists[i]:
                heapq.heappush(heap, (lists[i].val, i, lists[i]))

        while heap:
            node = heapq.heappop(heap)
            idx = node[1]
            result.next = node[2]

            result = result.next
            if result.next:
                heapq.heappush(heap, (result.next.val, idx, result.next))

        return root.next