Untitled

mail@pastecode.io avatar
unknown
c_cpp
2 years ago
1.4 kB
0
Indexable
Never
#include <bits/stdc++.h> 

vector<int> mergeKSortedArrays(vector<vector<int>>&kArrays, int k)
{
    // comparator lambda function to build min heap by comparing elements of two respective array.
    auto comparator = [&kArrays](pair<int, int> val1, pair<int, int> val2) {
        int arrayNo1 = val1.first;
        int eleIndx1 = val1.second;
        int arrayNo2 = val2.first;
        int eleIndx2 = val2.second;

        return kArrays[arrayNo1][eleIndx1] > kArrays[arrayNo2][eleIndx2];
    };

    // Each node of min Heap contains (arrayNo, index of element in that array)
    priority_queue<pair<int, int>, vector<pair<int, int> >, decltype(comparator)> minHeap(comparator);
    vector<int> mergedArr;

    // Initially insert 0th element of all the K arrays.
    for (int i = 0; i < k; i++)
        minHeap.push({i, 0});

    while(!minHeap.empty()) {
        // Extract min element among the arrays.
        int arrayNo = minHeap.top().first;
        int eleIndx = minHeap.top().second;
        minHeap.pop();

        mergedArr.push_back(kArrays[arrayNo][eleIndx]);

        // Move the index to next element for the array whose element was choosen as it was min of all.
        if (eleIndx + 1 < kArrays[arrayNo].size())
            minHeap.push({arrayNo, eleIndx + 1});
    }

    return mergedArr;
}