Programming/leetcode

23. Merge k Sorted Lists

홍열 2021. 2. 19. 13:57
728x90

ListNode를 가진 List들을 정렬해서 하나의 ListNode로 만드는 문제이다. 

 

두 가지 방법이 있다.

 

1) ListNode를 한번 씩 다 거쳐서, 새로운 list를 만들고, 다시 ListNode로 만드는 방법

 

    def mergeKLists(self, lists: List[ListNode]) -> ListNode:
        head = point = ListNode(0)
        node = []
        
        for item in lists:
            # 2차원 배열에서 하나씩 가져와서 [0]번의 ListNode부터 끝번까지 돌면서 val만 저장
            #it = item[0]
            while item:
                node.append(item.val)
                item = item.next
        
        for x in sorted(node):
            point.next = ListNode(x)
            point = point.next
        
        # head의 처음에는 0이 저장되어 있으므로, 그 다음것부터 보내도록 변경
        return head.next

하지만 이 방법은 시간이 O(n^2)이 걸린다. lists에서 각 ListNode를 돌면서 배열을 만들기때문이다.

 

이것을 해결하기 위해서는 lists에서 최솟값만 찾고, 다음으로 넘기는 식으로 해야된다. 

책에서는 이걸 min heap을 이용하여 풀었다. 

 

2) min heap을 사용하여서 각 ListNode중에서 최소값을 뽑는 방법

 

heap에 각 ListNode의 첫번째 값과 index와 ListNode를 저장한다.

index와 ListNode까지 저장하는 이유는 heap에 중복을 넣을 수 없기때문에 키를 다양화 하기 위함이다. 

또한 ListNode의 val이 같을 경우, 그 다음 값을 보고 어떻 값을 넣을지 결정한다. 

 

root와 result을 만들고, result를 옮겨가면서 최솟값을 저장해나가자.

    def mergeKLists(self, lists: List[ListNode]) -> ListNode:
        # list였다면 이렇게 headpify 사용해서 하면 되는데...
        # ListNode여서 하나씩 돌아가면서 일일히 추가해야됨
        # heap = heapq.heapify(lists)
        
        heap = []
        root = result = ListNode(None)
        for i, item in enumerate(lists):
            #item.val이 같으면 i를 비교
            if item:
                heapq.heappush(heap, (item.val, i, item))

        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