[leetcode] merge k sort list

/**
 * Definition for singly-linked list.
 * struct ListNode {
 *     int val;
 *     ListNode *next;
 *     ListNode(int x) : val(x), next(NULL) {}
 * };
 */
class Solution {
public:
    typedef struct{
        bool operator()(ListNode* node_1, ListNode* node_2){
            return node_1->val < node_2->val;
        }
    }Comp;
    ListNode *mergeKLists(vector<ListNode *> &lists) {
        // Start typing your C/C++ solution below
        // DO NOT write int main() function
        int k=0 ;
        if (lists.size()==0)
            return NULL;
        if (lists.size() == 1)
            return lists[0];
            
        ListNode* head = NULL;
        ListNode* pre = NULL;
        multiset<ListNode*, Comp> temp_heap;
        
        for(int i = 0; i< lists.size();i++){
            if ( lists[i] != NULL){
                temp_heap.insert(lists[i]);
                k++;
                lists[i] = lists[i]->next;
            }
        }
        
        while(k>1){
            ListNode* temp = *temp_heap.begin();
            if (head == NULL){
                head = temp;
            }
            else
                pre->next = temp;
            pre = temp;
            temp_heap.erase(temp_heap.begin());
            if (temp->next != NULL)
                temp_heap.insert(temp->next);
            else
                k--;
            
        }
        
        if(pre != NULL)
		pre->next =*temp_heap.begin();
	else
		head = *temp_heap.begin();
        return head;
    }
};

1. 依然用multiset實現heap

2. 

if(pre != NULL)
	pre->next =*temp_heap.begin();
else
	head = *temp_heap.begin();

最開始沒有做判斷，導致如果只有一個list爲非空是，出錯