[Java101] LeetCode: 23. Merge k Sorted Lists: Priority Queue

/**
 * 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; }
 * }
 */
class Solution {
    public ListNode mergeKLists(ListNode[] lists) {
        if (lists == null || lists.length == 0) {
            return null;
        }

        PriorityQueue<ListNode> pq = new PriorityQueue<>(
            lists.length, new Comparator<ListNode>() {
                @Override
                public int compare(ListNode node1, ListNode node2) {
                    if (node1.val < node2.val) {
                        return -1;
                    } else if (node1.val > node2.val) {
                        return 1;
                    } else {
                        return 0;
                    }
                }
            }
        );
        for (ListNode node: lists) {
            if (node != null) pq.add(node);
        }
        ListNode head = new ListNode();
        ListNode cur = head;
        while (!pq.isEmpty()) {
            ListNode polled = pq.poll();
            cur.next = polled;
            cur = polled;
            if (polled.next != null) {
                pq.add(polled.next);
            }
        }
        cur.next = null;
        return head.next;
    }
}