C++ O(NK log(NK)) Easy Solution with Comments, using min heap, Suggestions Welcome :)


#1
// A structure to store [element, {row , col}]
typedef pair<int, pair<int, int>> pi;

vector<int> Solution::solve(vector<vector<int>> &A) 
{
    // The given array A can be thought as
    // K arrays, where each array is present in a row
    vector<int> res;
    int R = A.size(), C = A[0].size(), row, col, ele;
    
    // Create a min heap of above typedef
    // Note that first element of first pair of typedef is key for heap
    priority_queue<pi, vector<pi>, greater<pi>> pq;
    
    // Traverse in Column major order
    // i.e. Add all first elements of k arrays to min heap
    for(int i = 0; i < R; i++)
    {
        // Push element, along with it's location
        pq.push({A[i][0], {i, 0} });
    }
    
    // Run a loop while PQ not empty
    while(!pq.empty())
    {
        // Since min element is at root
        // Get the element
        auto temp = pq.top();
        row = temp.second.first;
        col = temp.second.second;
        
        // Pop from heap
        pq.pop();
        
        // Add the popped element to our res
        res.emplace_back(temp.first);
        
        // Add next element from the array to which popped element belonged
        // But first, Check the next column element from popped ele location
        if(col + 1 < C)
        {
            pq.push({A[row][col + 1], {row, col + 1} });
        }
    }
    return res;
}