Leetcode 23: Merge k sorted lists

leetcode_23.cpp

#include <iostream>
#include <vector>
#include <map>

using namespace std;


//Definition for singly-linked list.
struct ListNode {
    int val;
    ListNode *next;
    ListNode(int x) : val(x), next(NULL) {}
};
 
class Solution {
public:
    ListNode* mergeKLists(vector<ListNode*>& lists) {
        if (lists.size()==0)
            return NULL;
        if (lists.size()==1)
            return lists[0];
        
        //pointers to current nodes as we traverse through the lists.
        vector<ListNode*> nodes; 
        //Initialise to head of each list
        for (int i=0; i<lists.size(); i++) {
            nodes.push_back(lists[i]);
            if (lists[i] != NULL)
                valsmap.insert({lists[i]->val, i});
        }
        
        auto min_pair = find_min(nodes, -1);
        int lidx_min0 = min_pair.first;  //smallest value of nodes set
        int lidx_min1 = min_pair.second; //second smalles value
        
        if (lidx_min0 == -1 && lidx_min1 == -1)
            return NULL;
        else if (lidx_min1 == -1)
            return nodes[lidx_min0];
        
        ListNode* head = nodes[lidx_min0];
        
        while (lidx_min1 >= 0) {
            //see if we can switch to next node in same list
            if (nodes[lidx_min0]->next && 
                nodes[lidx_min0]->next->val <= nodes[lidx_min1]->val) {
                int oldval = nodes[lidx_min0]->val;
                nodes[lidx_min0] = nodes[lidx_min0]->next;

                valsmap_erase(oldval, lidx_min0);
                valsmap.insert({nodes[lidx_min0]->val, lidx_min0});
            }
            else { //switch lidx_min0 to point to lidx_min1 list node
                int oldval = nodes[lidx_min0]->val;
                ListNode* next_node = nodes[lidx_min0]->next;
                nodes[lidx_min0]->next = nodes[lidx_min1];
                nodes[lidx_min0] = next_node;

                valsmap_erase(oldval, lidx_min0);
                if (next_node)
                    valsmap.insert({next_node->val, lidx_min0});
                
                min_pair = find_min(nodes, lidx_min1);
                lidx_min0 = min_pair.first; 
                lidx_min1 = min_pair.second; 
            }
        }
        return head;
    }
private:
    pair<int, int> find_min_brute(vector<ListNode*>& nodes, int default_idx) {
        int idx_min0 = -1, idx_min1 = -1;
        for (int i=0; i<nodes.size(); i++) {
            if (nodes[i]==NULL)
                continue;
            if (idx_min0 == -1) 
                idx_min0 = i;
            else if (nodes[idx_min0]->val > nodes[i]->val) {
                if (idx_min1==-1 || nodes[idx_min1]->val > nodes[idx_min0]->val)
                    idx_min1 = idx_min0;
                idx_min0 = i;
            }
            else if (idx_min1 == -1)
                idx_min1 = i;
            else if (nodes[idx_min1]->val > nodes[i]->val)
                idx_min1 = i;
        }
        //smallest needs to be the previous second smallest.
        //our find algorithm above can return different order if 
        //the values are the same, so let's rearrange here if we 
        //need to.
        if (default_idx != -1 && idx_min0 != default_idx) {
            idx_min1 = idx_min0;
            idx_min0 = default_idx;
        }
        return {idx_min0, idx_min1};
    }

    multimap<int, int> valsmap; //key=nodeval; value=listindex

    pair<int, int> find_min(vector<ListNode*>& nodes, int default_idx) {
        auto smallest = valsmap.begin();
        if (smallest == valsmap.end())
            return {-1, -1};
        if (default_idx != -1 && smallest->second != default_idx)
            return { default_idx, smallest->second };
        int idx_min0 = smallest->second;
        smallest++;
        int idx_min1 = (smallest == valsmap.end()) ? -1 : smallest->second;
        return { idx_min0, idx_min1 };
    }

    void valsmap_erase(int key, int val) {
        auto range = valsmap.equal_range(key);
        for (auto i = range.first; i != range.second; ++i) {
            if (i->second == val) {
                valsmap.erase(i);
                return;
            }
        }
    }
};

vector<ListNode*> nodes_owner;  //so we can easily delete the nodes.

void delete_nodes() {
    while (nodes_owner.size() > 0) {
        ListNode* node = nodes_owner.back();
        delete node;
        nodes_owner.pop_back();
    }
}

ListNode* add_node(ListNode* node, int new_val) {
    ListNode* new_node = new ListNode(new_val);
    nodes_owner.push_back(new_node); //memory management
    if (node)
        node->next = new_node;
    return new_node;
}

vector<ListNode*> initLists() {
    vector<ListNode*> lists;
    
/*    ListNode* n1 = add_node(NULL, -1);
    lists.push_back(n1);
    n1 = add_node(n1, 1);
    
    ListNode* n2 = add_node(NULL, -3);
    lists.push_back(n2);
    n2 = add_node(n2, 1);
    n2 = add_node(n2, 4);
    
    ListNode* n3 = add_node(NULL, -2);
    lists.push_back(n3);
    n3 = add_node(n3, -1);
    n3 = add_node(n3, 0);
    n3 = add_node(n3, 2);
*/
    ListNode* n1 = add_node(NULL, 1);
    lists.push_back(n1);
    n1 = add_node(n1, 4);
    n1 = add_node(n1, 5);
   
    ListNode* n2 = add_node(NULL, 1);
    lists.push_back(n2);
    n2 = add_node(n2, 3);
    n2 = add_node(n2, 4);
    
    ListNode* n3 = add_node(NULL, 2);
    lists.push_back(n3);
    n3 = add_node(n3, 6);

    return lists;
}

void print_list(ListNode* node) {
    while (node) {
        cout << node->val;
        if (node->next)
            cout << " -> ";
        node = node->next;
    }
    cout << endl;
}

int main() {
    vector<ListNode*> lists = initLists();
    Solution s;
    ListNode* node = s.mergeKLists(lists);
    print_list(node);
    delete_nodes();
    return 0;
}