Merging K Sorted Lists - Amazon Top Interview Questions

                        Solution in C++ :

vector<int> merge(vector<int>& a, vector<int>& b) {
    vector<int> temp;
    int i = 0, j = 0;
    if (a.size() == 0) return b;
    if (b.size() == 0) return a;

    while (i < a.size() && j < b.size()) {
        if (a[i] <= b[j]) {
            temp.push_back(a[i]);
            i++;
        } else {
            temp.push_back(b[j]);
            j++;
        }
    }
    while (i < a.size()) {
        temp.push_back(a[i]);
        i++;
    }
    while (j < b.size()) {
        temp.push_back(b[j]);
        j++;
    }

    return temp;
}

vector<int> helper(vector<vector<int>>& lists, int st, int end) {
    if (st == end) return lists[st];
    int mid = st + (end - st) / 2;
    vector<int> a = helper(lists, st, mid);
    vector<int> b = helper(lists, mid + 1, end);
    return merge(a, b);
}

vector<int> solve(vector<vector<int>>& lists) {
    int k = lists.size();
    if (k == 0) return {};

    return helper(lists, 0, k - 1);
}
                    

                        Solution in Python : 
                            
class Solution:
    def solve(self, lists):
        min_heap = []

        for i in range(len(lists)):
            if len(lists[i]) > 0:
                heapq.heappush(min_heap, (lists[i][0], i, 0))

        res = []

        while len(min_heap) > 0:
            a, i, j = heapq.heappop(min_heap)

            if j + 1 < len(lists[i]):
                heapq.heappush(min_heap, (lists[i][j + 1], i, j + 1))

            res.append(a)

        return res