sorted matrix 系列 378

Given an n x n matrix where each of the rows and columns is sorted in ascending order, return the kth smallest element in the matrix.

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

You must find a solution with a memory complexity better than O(n2).

Example 1:

Input: matrix = [[1,5,9],[10,11,13],[12,13,15]], k = 8
Output: 13
Explanation: The elements in the matrix are [1,5,9,10,11,12,13,13,15], and the 8th smallest number is 13

Example 2:

Input: matrix = [[-5]], k = 1
Output: -5

Constraints:

• n == matrix.length == matrix[i].length
• 1 <= n <= 300
• -109 <= matrix[i][j] <= 109
• All the rows and columns of matrix are guaranteed to be sorted in non-decreasing order.
• 1 <= k <= n2

Follow up:

• Could you solve the problem with a constant memory (i.e., O(1) memory complexity)?
• Could you solve the problem in O(n) time complexity? The solution may be too advanced for an interview but you may find reading this paper fun.

解法一：

使用min heap存储每行的头元素，逐个遍历取出k个元素，

将该元素的右侧下一个元素压入heap，最终得到第k个元素

时间复杂度：O(Klog(N))

class Solution { 
    public int kthSmallest(int[][] matrix, int k) {
        int m = matrix.length, n = matrix[0].length, ans = -1; // For general, the matrix need not be a square
        PriorityQueue<int[]> minHeap = new PriorityQueue<>(Comparator.comparingInt(o -> o[0]));
        // 每行第一个元素放入heap
        for (int r = 0; r < Math.min(m, k); ++r)
            minHeap.offer(new int[]{matrix[r][0], r, 0});
        // 开始从heap中取出元素，数k个
        for (int i = 1; i <= k; ++i) {
            int[] top = minHeap.poll();
            // 当前元素的位置
            int x = top[1], y = top[2];
            ans = top[0];
            if (y + 1 < n) minHeap.offer(new int[]{matrix[x][y + 1], x, y + 1});
        }
        return ans;
    }
}

解法二：二分猜答案

时间复杂度：O(N) 

class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        // 给定最大/最小值，二分猜答案
        long left = Integer.MIN_VALUE, right = Integer.MAX_VALUE;
        int result = 0;
        while(left <= right) {
            long mid = left + (right - left) / 2;
            // 如果比mid的元素小的个数少于k，那么需要往大了猜
            if(countSmallerOrEqual(matrix, mid) < k) {
                left = mid + 1;
            }
            // 如果比mid元素小的个数多于或者等于k，那么记录下来result，继续往小了猜
            else{
                result = (int)mid;
                right = mid - 1;
            }
        }
        return result;
    }
    private int countSmallerOrEqual(int[][] matrix, long mid) {
        int m = matrix.length, n = matrix[0].length;
        int result = 0;
        // 从右上角开始数
        int x = 0, y = n - 1;
        while(x < m && y >= 0) {
            // 如果比mid大，向左移动
            if(matrix[x][y] > mid) y--;
            // 如果比mid小，向下移动，并且记录当前行小于mid的元素
            else {
                result += y + 1;
                x++;
            }
        }
        return result;
    }
}