# Definition for singly-linked list.
# class ListNode:
#     def __init__(self, val=0, next=None):
#         self.val = val
#         self.next = next
from heapq import heappush, heappop
class Solution:
    '''
    Time Complexity: O(n log k)
    Space Complexity: O(k)

    where, k is number of linked lists, and n is total number of nodes across all lists

    NOTE:
    When using heapq with custom classes/objects, use a tuple of the form to be stored:
        (key, unique_id, object)

    Reason:
    - heapq does not support a custom comparator or key function.
    - It compares tuple elements in order.
    - First, it compares `key`.
    - If keys are equal, it compares `unique_id`.
    - Without a unique_id, Python would attempt to compare the objects
    themselves, which raises a TypeError since class instances cannot be compared by default.
    '''
    def mergeKLists(self, lists: List[Optional[ListNode]]) -> Optional[ListNode]:
        heap = []
        dummy = ListNode(0)
        node = dummy

        lists = [head for head in lists if head]
        counter = 0
        for l in lists:
            if l:
                heappush(heap, (l.val, counter, l))
                counter += 1

        while heap:
            val, _, node.next = heappop(heap)
            node = node.next
            if node.next:
                heappush(heap, (node.next.val, counter, node.next))
                counter += 1
        
        return dummy.next

// // Priority Queue (As quick as Divide and Conquer, theoretical time complexity is best)
/**
 * Definition for singly-linked list.
 * public class ListNode {
 *     int val;
 *     ListNode next;
 *     ListNode() {}
 *     ListNode(int val) { this.val = val; }
 *     ListNode(int val, ListNode next) { this.val = val; this.next = next; }
 * }
 */
public ListNode mergeKLists(ListNode[] lists) {
    ListNode dummy = new ListNode();
    ListNode current = dummy;
    Queue<ListNode> pq = new PriorityQueue<>(new Comparator<ListNode>() {
            @Override
            public int compare(ListNode l1, ListNode l2) {
                return l1.val - l2.val;
            }
        });

    for (ListNode list: lists) {
        if (list != null) {
            pq.add(list);
        }
    }

    while (pq.size() > 1) {
        current.next = pq.poll();
        current = current.next;
        if (current.next != null) {
            pq.add(current.next);
        }
    }

    current.next = pq.poll();

    return dummy.next;
}