Two layer binary search

378. Kth Smallest Element in a Sorted Matrix

We want to use the lower_bound to find a range of data qualify ≥ condition, based on this, we can make sure lo always exist in the array.

Example 1

For example

1, 2, 2, 2, 4, 5

If we want to find the 3th largest element:

step by step will be

lo = 1, hi = 5

mid = (1 + 5 +1) / 2 = 3

we have 2 elements ≥ 3, so we need reduce the mid value to find more elements.

hi = mid-1=2

mid = (1 + 2 + 1) / 2 = 2

we have 5 elements ≥ 2, so we need keep the mid value as this might be a qualified case as we allow duplicated elements in the array.

Let’e move forward

lo = mid = 2

lo = 2 and hi = 2, we break, so 2 is the correct answer. We can see, by using this way, we can always find the correct answer.

We can also go over another example.

Example 2

1, 2, 3, 4, 5, 6

If we want to find the 3th largest element:

step by step will be

lo = 1, hi = 6

mid = (1 + 6+1) / 2 = 4

we have 3 elements ≥ 4, we update the low is equal to mid.

lo = 4, hi = 6

mid = (4 + 6+ 1) / 2 = 5

we have 2elements ≥ 5, we update the hi = mid -1

lo = 4, hi = 4

we find the correct result.

class Solution {
public:
    int kthSmallest(vector<vector<int>>& matrix, int k) {
        if (matrix.empty() || matrix[0].empty()) return INT_MIN;
        const int n = matrix.size(), m = matrix[0].size();
        k = n*m - k + 1;
        int lo = matrix[0][0], hi = matrix[n-1][m-1];
        while (lo < hi){
            int mid = (lo + hi + 1) / 2;
            int count = 0;
            for (int i = 0; i < n; ++i){
                auto& a = matrix[i];
                count += a.end() - lower_bound(a.begin(), a.end(), mid);
            }
            if (count >= k) lo = mid;
            else hi = mid - 1;
        }
        return lo;
    }
};