23. Merge k Sorted Lists

kukudas·2022년 3월 24일
0

Algorithm

목록 보기
24/46

여기서 그냥 아래처럼 해버리면
heapq에서 내부적으로 <를 사용해서 대소 비교를 하기에 ListNode() 객체는 <를 못쓰니 에러가 나옴.

heapq.heapush(h, (lists[i].val, lists[i]))

따라서 __lt__ 를 추가해 주거나

class ListNode:
    def __init__(self, x):
        self.val = x
        self.next = None

    def __lt__(self, other):
        return self.val < other.val

여기처럼 그냥 인덱스로 추가 비교 할 수 있도록해주면 문제없음.
비교할때 set에서 앞에 있는 순서대로 비교하니까 맨 앞에 val이 들어와야하고 그 다음에 인덱스 그리고 마지막에 노드가 들어가도록 해줘야함.

class Solution:
    def mergeKLists(self, lists):
        head = cur = ListNode(None)
        h = []

        # 각 연결리스트의 root를 힙에 넣어줌
        for i in range(len(lists)):
            # [[]] 같은거 들어올 수 있음
            if lists[i]:
                heapq.heappush(h, (lists[i].val, i, lists[i]))

        # 루트를 힙에서 빼고 루트 다음을 힙에 넣어줌
        while h:
            pop = heapq.heappop(h)
            cur.next = pop[2]
            cur = cur.next

            # 뺀 거 다음이 있으면 다시 힙에 넣어줌
            if pop[2].next:
                heapq.heappush(h, (pop[2].next.val, pop[1], pop[2].next))

        return head.next

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

0개의 댓글