LeetCode: Merge k Sorted Lists

Merge k sorted linked lists and return it as one sorted list. Analyze and describe its complexity.

---

I have two ideas for this problem. One is divide and conquer, merge every two linked lists recursively. The other one is use a heap to store head nodes of each linked list, and each node pop the smallest one and push its successor into the current heap until there is nothing in the heap. I implemented the first approach in C++ and the second one in Python.

C++ Code:

/*
 * func: merge_two_lists
 * goal: merge 2 linked lists into one linked list
 * @param list1: linked list 1
 * @param list2: linked list 2
 * return: the new linked list
 * complexity: time O(n), space O(1)
 */
ListNode* merge_two_lists(ListNode* list1, ListNode* list2){
    ListNode* iter1 = list1;
    ListNode* iter2 = list2;
    ListNode* result_head = nullptr;
    ListNode* result_iter = new ListNode(0);
    //Merge two lists
    while(iter1 != nullptr && iter2 != nullptr){
        if(iter1->val <= iter2->val){
            result_iter->next = iter1;
            iter1 = iter1->next;
        }else{
            result_iter->next = iter2;
            iter2 = iter2->next;
        }
        //Get the head of new list
        if(result_head == nullptr){
            result_head = result_iter->next;
        }
        result_iter = result_iter->next;
    }
    //Concatenate remaining of list1
    while(iter1 != nullptr){
        result_iter->next = iter1;
        iter1 = iter1->next;
        if(result_head == nullptr){
            result_head = result_iter->next;
        }
        result_iter = result_iter->next;
    }
    //Concatenate remaining of list1
    while(iter2 != nullptr){
        result_iter->next = iter2;
        iter2 = iter2->next;
        if(result_head == nullptr){
            result_head = result_iter->next;
        }
        result_iter = result_iter->next;
    }
    
    return result_head;
}

/*
 * func: merge_k_lists_helper
 * goal: recursive function to help merge lists
 * @param lists: vector of head node
 * @param begin: index of list 1
 * @param end: index of list 2
 * return: merged list
 */
ListNode* merge_k_lists_helper(vector<ListNode*> lists, int begin, int end){
    if(begin > end){
        return nullptr;
    }
    if(begin == end){
        return lists.at(begin);
    }
    if(begin + 1 == end){
        return merge_two_lists(lists.at(begin), lists.at(end));
    }
    int mid = begin + (end-begin)/2;
    ListNode* left = merge_k_lists_helper(lists, begin, mid);
    ListNode* right = merge_k_lists_helper(lists, mid+1, end);
    return merge_two_lists(left, right);
}


/*
 * func: merge_k_lists
 * goal: merge k linked list into one linked list
 * @param lists: vector of head node
 * return: the new linked list
 */
/*
 * divide and conquer
 * complexity: time O(nlogk), space O(1)
 */
ListNode* merge_k_lists(vector<ListNode*> lists){
    return merge_k_lists_helper(lists, 0, lists.size()-1);
}

Python Code:

# func: merge k linked list into one list
# @param lists: lists of head nodes
# @return: head node of merged list
def merge_k_lists(lists):
    if not lists:
        return None
    if len(lists) == 1:
        return lists[0]
    head = None
    curr_node = None

    heap = []
    for node in lists:
        if node is not None:
            heapq.heappush(heap, (node.val, node))

    #Use heap to construct
    while heap:
        value, min_node = heapq.heappop(heap)
        if head is None:
            head = min_node

        if curr_node is None:
            curr_node = min_node
        else:
            curr_node.next = min_node
            curr_node = curr_node.next

        min_node = min_node.next

        if min_node is not None:
            heapq.heappush(heap, (min_node.val, min_node))

    return head