LeetCode/cpp/kth-smallest-element-in-a-sorted-matrix.cpp

#include <algorithm>
#include <cassert>
#include <iostream>
#include <vector>

using namespace std;

class Solution {
public:
    int kthSmallest(vector<vector<int>>& matrix, int k)
    {
        int low = matrix.front().front();
        int high = matrix.back().back();

        while (low < high) {
            int mid = low + (high - low) / 2;

            int rank = 0;
            for (const auto& row : matrix) {
                rank += upper_bound(row.begin(), row.end(), mid) - row.begin();
            }

            if (rank < k) {
                low = mid + 1;
            } else {
                high = mid;
            }
        }

        return low;
    }
};

int main()
{
    Solution s;

    vector<vector<int>> m;
    int k;

    m = {
        { 1, 5, 9 },
        { 10, 11, 13 },
        { 12, 13, 15 }
    };
    k = 8;
    assert(s.kthSmallest(m, k) == 13);

    m = { { -5 } };
    k = 1;
    assert(s.kthSmallest(m, k) == -5);
}