K-way Merge

kthSmallest.cpp

// min-heap
int kthSmallest(vector<vector<int>>& matrix, int k) {
    typedef pair<int, pair<int,int>> ppi;
    priority_queue<ppi, vector<ppi>, greater<ppi>> pq; //greater<ppi> will sort by ppi.first in ascending order
    for(int i = 0; i < matrix.size(); i++){
        pq.push(make_pair(matrix[i][0], make_pair(i, 0)));
        // or pq.push( {matrix[i][0], {i, 0} });
    }
    int i = 0, j = 0;
    while(!pq.empty() && k!=0){
        ppi ele = pq.top();
        pq.pop();
        k--;
        i = ele.second.first;
        j = ele.second.second;
        if(j != matrix[i].size() - 1) //if current element is not the last element of current row
            pq.push(make_pair(matrix[i][j+1], make_pair(i, j+1)));
    }
    return matrix[i][j];
}

// max-heap
int kthSmallest(vector<vector<int>>& matrix, int k) {
    priority_queue<int> q;
    for (int i = 0; i < matrix.size(); ++i) {
        for (int j = 0; j < matrix[i].size(); ++j) {
            q.push(matrix[i][j]);
            if (q.size() > k) q.pop();
        }
    }
    return q.top();
}