Kth Smallest Element in a Sorted Matrix

378.cpp

class MinHeap {
private:
    // value -> (row, col)
    vector<tuple<int, int, int>> nodes;
    
    size_t used = 0;
public:
    MinHeap(size_t n) {
        nodes.resize(n);        
    }

    void insert(int val, int r, int c) {
        if (used++ == nodes.capacity()) {
            nodes.resize(nodes.capacity() * 2);
        }
        nodes[used - 1] = make_tuple(val, r, c);
        
        int k = used - 1;
        while (k > 0 && get<0>(nodes[k]) < get<0>(nodes[(k - 1) / 2])) {
            nodes[k].swap(nodes[(k - 1)/2]);
            k = (k - 1) / 2;
        }
    }
    
    tuple<int, int, int>& peak() {
        return nodes[0];
    }
    
    void remove_min() {
        nodes[0].swap(nodes[used - 1]);
        used--;
        
        for (int k = used / 2 - 1; k >= 0; k--) {
            if (get<0>(nodes[k]) > get<0>(nodes[k * 2 + 1]))
                nodes[k].swap(nodes[k * 2 + 1]);
            
            if (k * 2 + 2 <= (used - 1) && get<0>(nodes[k]) > get<0>(nodes[k * 2 + 2])) {
                nodes[k].swap(nodes[k * 2 + 2]);               
            }
        }
    }
};

class Solution {
public:
    int kthSmallest(vector<vector<int>>& matrix, int k) {
        int row_count = min((int)matrix.size(), k);
        MinHeap h(row_count);
        for (int i = 0; i < row_count; i++) {
            h.insert(matrix[i][0], i, 0);
        }
        
        int idx = 1;
        int ans;
        do {
            auto tu = h.peak();
            h.remove_min();
            ans = get<0>(tu);
            int row = get<1>(tu);
            int col = get<2>(tu);

            if (++col < matrix[0].size()) {
                h.insert(matrix[row][col], row, col);   
            }
        } while (idx++ != k);
        
        return ans;
    }
};