378. Kth Smallest Element in a Sorted Matrix - Medium

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 ≤ n2.

 

M1: heap，建立一个min heap，先把最小元素{0, 0}放进heap。用cnt计数，如果++cnt == k时，找到第k小的元素，返回。

因为数组每行每列都是有序的，每次循环时把heap中的最小值poll出来，为了避免重复，如果是第一行，把该元素右边的元素加入heap，然后把该元素下方的元素加入heap。用一个参数cnt计数，当++cnt = k里，返回matrix对应坐标的值。

time: O(klogk), space: O(k) 

class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        PriorityQueue<int[]> minHeap = new PriorityQueue<>(k, (a, b) -> matrix[a[0]][a[1]] - matrix[b[0]][b[1]]);
        int cnt = 0;
        minHeap.add(new int[] {0, 0});
        while(true) {
            int[] pair = minHeap.poll();
            if(++cnt == k)
                return matrix[pair[0]][pair[1]];
            
            if(pair[0] == 0 && pair[1] + 1 < matrix[0].length)
                minHeap.add(new int[] {pair[0], pair[1] + 1});
                
            if(pair[0] + 1 < matrix.length)
                minHeap.add(new int[] {pair[0] + 1, pair[1]});                
        }
    }
}

另一种写法：

从最小的元素开始，依次访问最小元素的neighbor（当neighbor下标不越界且没被访问过，入队，并标记为已访问），循环k次。最后最小堆的堆顶元素就是第k小元素的坐标

time: O(klogk), space: O(k + n^2)

class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        boolean[][] visited = new boolean[matrix.length][matrix[0].length];
        PriorityQueue<int[]> minHeap = new PriorityQueue<>((a, b) -> matrix[a[0]][a[1]] - matrix[b[0]][b[1]]);
        minHeap.offer(new int[] {0, 0});
        
        for(int i = 0; i < k - 1; i++) {
            int[] idx = minHeap.poll();
            int row = idx[0], col = idx[1];
            visited[row][col] = true;
            
            if(row + 1 < matrix.length && !visited[row+1][col]) {
                minHeap.offer(new int[] {row + 1, col});
                visited[row+1][col] = true;
            }
            if(col + 1 < matrix[0].length && !visited[row][col+1]) {
                minHeap.offer(new int[] {row, col + 1});
                visited[row][col+1] = true;
            }
        }
        
        int[] res = minHeap.peek();
        return matrix[res[0]][res[1]];
    }
}

 

M2: binary search，原理同 287. find the duplicate number https://www.cnblogs.com/fatttcat/p/9986854.html

在number value的基础上进行binary search，范围是min~max。每次计算matrix中有多少个<= mid值的数并缩小搜索范围：如果cnt < k，说明只有不到k个数小于等于mid，第k个数应该在mid之后的部分，left=mid+1；反之如果cnt >= k，说明第k个数在mid之前，right=mid-1

time: O(Nlog(max - min)), space: O(1)

class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        int n = matrix.length;
        int left = matrix[0][0], right = matrix[n - 1][n - 1];
        while(left <= right) {
            int mid = left + (right - left) / 2;
            int cnt = lessEqualK(matrix, mid);
            if(cnt < k) {
                left = mid + 1;
            } else {
                right = mid - 1;
            }
        }
        return left;
    }
    
    public int lessEqualK(int[][] matrix, int val) {
        int i = matrix.length - 1, j = 0;
        int cnt = 0;
        while(i >= 0 && j <= matrix.length - 1) {
            if(matrix[i][j] > val) {
                i--;
            } else {
                cnt += i + 1;
                j++;
            }
        }
        return cnt;
    }
}