[LeetCode 378] Kth Smallest Element in a Sorted Matrix

Given a n x n matrix where each of the rows and columns are sorted in ascending order, find the kth smallest element in the matrix.

Note that it is the kth smallest element in the sorted order, not the kth distinct element.

Example:

matrix = [
   [ 1,  5,  9],
   [10, 11, 13],
   [12, 13, 15]
],
k = 8,

return 13.

Note:
You may assume k is always valid, 1 ≤ k ≤ n^2.

The key observation here is that we have sorted rows and columns separately. The entire matrix is not necessrily sorted top to bottom, left to right. The given example is misleading on purpose.

Solution 1. Min heap, O(k * log n) runtime, O(n) space.

1. Add the 1st row to a min heap with the value, row and col info.

2. poll the smallest value out of min heap then add its next value on the same column. Repeat this k times. Because each row is sorted and each column is also sorted, it is guranteed that we always poll the next smallest value in the entire matrix.

class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        int m = matrix.length, n = matrix[0].length;
        PriorityQueue<int[]> minPq = new PriorityQueue<>(Comparator.comparingInt(a -> a[0]));
        for(int i = 0; i < n; i++) {
            minPq.add(new int[]{matrix[0][i], 0, i});
        }
        
        int ans = 0;
        while(k > 0) {
            int[] curr = minPq.poll();
            ans = curr[0];
            if(curr[1] < matrix.length - 1) {
                minPq.add(new int[]{matrix[curr[1] + 1][curr[2]], curr[1] + 1, curr[2]});
            }
            k--;
        }
        return ans;
    }
}

Variation of Solution 1.

In the above solution, we add the entire 1st row to the min heap initially. We can instead choose only to add matrix[0][0] and then for each element that is polled from the min heap, we try to add the number that is just to its right and the number that is just to its bottom. However, we will need O(N^2) extra book keeping to track if we have already added a number.

Consider this example: when 3 is polled, we try to add 5 and 7; when 6 is polled, we try to add 7 and 11. If there is no tracking, 7 will be added twice!

1 3 5

6 7 12

11 14 14

class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        PriorityQueue<int[]> minPq = new PriorityQueue<>(Comparator.comparingInt(a -> a[0]));
        boolean[][] added = new boolean[matrix.length][matrix[0].length];
        minPq.add(new int[]{matrix[0][0], 0, 0});
        added[0][0] = true;
        int ans = 0;
        while(k > 0) {
            int[] curr = minPq.poll();
            ans = curr[0];
            if(curr[1] < matrix.length - 1 && !added[curr[1] + 1][curr[2]]) {
                minPq.add(new int[]{matrix[curr[1] + 1][curr[2]], curr[1] + 1, curr[2]});
                added[curr[1] + 1][curr[2]] = true;
            }
            if(curr[2] < matrix[0].length - 1 && !added[curr[1]][curr[2] + 1]) {
                minPq.add(new int[]{matrix[curr[1]][curr[2] + 1], curr[1], curr[2] + 1});
                added[curr[1]][curr[2] + 1] = true;
            }
            k--;
        }
        return ans;
    }
}

Solution 2. Binary Search on possible answer range, O(maxV * N) runtime, O(1) space.

Each time we check the middle value MID of the current search range: if there are < k numbers that are <= MID, search on [mid + 1, right]; else search on[left, mid].

The countSmallerAndEqual method takes O(N + N) which is O(N) runtime. This is because for matrix[i][j], we have matrix[k][j] >= matrix[i][j], k > i; matrix[i][k] >= matrix[i][j], k > j.

So as we search for the total number of values that are <= MID, j can only decrease as i increases, we'll never reset j back to matrix[0].length - 1. Essentially, this count method just increases i from 0 to n - 1 and decreases j from n - 1 to 0.

class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        int l = matrix[0][0], r = matrix[matrix.length - 1][matrix[0].length - 1];
        
        while(l < r - 1) {
            int mid = l + (r - l) / 2;
            if(countSmallerAndEqual(matrix, mid) < k) {
                l = mid + 1;
            }
            else {
                r = mid;
            }
        }
        if(countSmallerAndEqual(matrix, l) >= k) {
            return l;
        }
        return r;
    }
    private int countSmallerAndEqual(int[][] matrix, int t) {
        int cnt = 0, j = matrix[0].length - 1;
        for(int i = 0; i < matrix.length; i++) {
            while(j >= 0 && matrix[i][j] > t) {
                j--;
            }
            cnt += j + 1;
        }
        return cnt;
    }
}

[LeetCode 378] Kth Smallest Element in a Sorted Matrix

猜你喜欢