23. Merge k Sorted Lists

• categories: Code, Interview Question, leetcode, Hard
• source: https://leetcode.com/problems/merge-k-sorted-lists
• topics: Linked List, Heap

Description

You are given an array of k linked-lists lists, each linked-list is sorted in ascending order.

Merge all the linked-lists into one sorted linked-list and return it.

Idea

• Use heap

Code

class Solution:
    def mergeKLists(self, lists: List[Optional[ListNode]]) -> Optional[ListNode]:
        heap = []
        head = tail = ListNode(0) # Dummies to handle edge cases
        
        # Push initial nodes of all non-empty lists into the heap
        for i, node in enumerate(lists):
            if node:
                heappush(heap, (node.val, i, node))
        
        # Merge lists by extracting the smallest node from the heap
        while heap:
            _, i, node = heappop(heap)
            tail.next = node
            tail = tail.next
            if node.next:
                heappush(heap, (node.next.val, i, node.next))
        
        return head.next