Merge k Sorted Lists

 
// C++
/**
 * Definition for singly-linked list.
 * struct ListNode {
 *     int val;
 *     ListNode *next;
 *     ListNode(int x) : val(x), next(NULL) {}
 * };
 */
class Solution {
typedef pair<int, int> PairInt;
public:
    ListNode *mergeKLists(vector<ListNode *> &lists) {
        priority_queue<PairInt, vector<PairInt>, greater<PairInt> > q;
        // initialize heap with first element from each list
        int k = lists.size();
        for (int i = 0; i < k; i++) {
            if (lists[i]) {
                q.push(PairInt(lists[i]->val, i));
            }
        }

        ListNode tmp(0);
        ListNode *p = &tmp;
       
        while (!q.empty()) {
            int i = q.top().second;
            q.pop();
            p->next = lists[i];
            p =p->next;
            lists[i] = lists[i]->next;
           
            if (lists[i]) {
                q.push(PairInt(lists[i]->val, i));
            }
        }
       
        return tmp.next;
    }
};


// Java
 public ListNode mergeKLists(ArrayList<ListNode> lists) {
        if (lists == null || lists.size() == 0) {  // this was missed.
            return null;
        }
        // PriorityQueue is the heap implementation in Java
        PriorityQueue<ListNode> pq = new PriorityQueue<ListNode>(lists.size(), new Comparator<ListNode>() {
            public int compare(ListNode l1, ListNode l2) {
                return l1.val - l2.val;
            }
        });
        for (ListNode head : lists) {
            if (head != null) {  // this was missed.
                pq.add(head);
            }
        }
        ListNode head = new ListNode(0);
        ListNode p = head;
        while (pq.size() > 0) {
            p.next = pq.poll();
            p = p.next; // this was missed.
            if (p.next != null) {
                pq.add(p.next);
            }
        }
        return head.next;
    }