【Day 89 】2023-01-28 - 378. 有序矩阵中第 K 小的元素

[azl397985856]

378. 有序矩阵中第 K 小的元素

入选理由

暂无

题目地址

https://leetcode-cn.com/problems/kth-smallest-element-in-a-sorted-matrix/

前置知识

• 二分查找
• 堆

题目描述

给定一个 n x n 矩阵，其中每行和每列元素均按升序排序，找到矩阵中第 k 小的元素。
请注意，它是排序后的第 k 小元素，而不是第 k 个不同的元素。

 

示例：

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

返回 13。
 

提示：
你可以假设 k 的值永远是有效的，1 ≤ k ≤ n2 。

[whoam-challenge]

class Solution:
def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
n = len(matrix)
pq = [(matrix[i][0], i, 0) for i in range(n)]
heapq.heapify(pq)

    ret = 0
    for i in range(k - 1):
        num, x, y = heapq.heappop(pq)
        if y != n - 1:
            heapq.heappush(pq, (matrix[x][y + 1], x, y + 1))
    
    return heapq.heappop(pq)[0]

[Abby-xu]

class Solution:
def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
m, n = len(matrix), len(matrix[0]) # For general, the matrix need not be a square
minHeap = [] # val, r, c
for r in range(min(k, m)):
heappush(minHeap, (matrix[r][0], r, 0))

    ans = -1  # any dummy value
    for i in range(k):
        ans, r, c = heappop(minHeap)
        if c+1 < n: heappush(minHeap, (matrix[r][c + 1], r, c + 1))
    return ans

[zywang0]

class Solution { // 14 ms, faster than 55.67%
    public int kthSmallest(int[][] matrix, int k) {
        int m = matrix.length, n = matrix[0].length; // For general, the matrix need not be a square
        PriorityQueue<Integer> maxHeap = new PriorityQueue<>((o1, o2) -> Integer.compare(o2, o1));
        for (int r = 0; r < m; ++r) {
            for (int c = 0; c < n; ++c) {
                maxHeap.offer(matrix[r][c]);
                if (maxHeap.size() > k) maxHeap.poll();
            }
        }
        return maxHeap.poll();
    }
}

[Ryanbaiyansong]

class Solution:
    def kthSmallest(self, a: List[List[int]], k: int) -> int:
        m, n = len(a), len(a[0])
        # n * n
        # 二分搜索找第k个

        l, r = min(x for row in a for x in row), max(x for row in a for x in row)
        print(l, r)
        
        def check(x: int) -> bool:
            '''
            check how many values is smaller than value
            '''
            i = 0
            j = n - 1
            cnt = 0
            while i < n and j >= 0:
                if a[i][j] <= x:
                    cnt += (j + 1)
                    i += 1
                else:
                    j -= 1
            return cnt >= k

        while l < r:
            mid = (l + r) // 2
            if check(mid):
                r = mid 
            else:
                l = mid + 1

        return l

[chenming-cao]

解题思路

堆。多路归并问题。类似合并k个已经排好序的链表。建立小顶堆，在堆中储存元素的row和column下标数组。首先将所有n行的第一个元素的坐标放入堆中，之后每次从堆中取出最小值的坐标，然后保持row不变，column增加1。将更新的坐标放入堆中。以此方法重复k次，第k次从堆中取出的坐标即为第k小元素的坐标，根据坐标返回值即可。

代码

class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        int n = matrix.length;
        // use heap to store the row and column indices
        PriorityQueue<int[]> heap = new PriorityQueue<>((a, b) -> matrix[a[0]][a[1]] - matrix[b[0]][b[1]]);
        // put the first element from each row to the heap
        for (int i = 0; i < n; i++) {
            heap.offer(new int[]{i, 0});
        }
        // sort to get 1st to (k - 1)th smallest elements
        for (int j = 0; j < k - 1; j++) {
            int[] cur = heap.poll();
            // check column index to see whether it is within the range, if yes, put the next element from the current row to the heap
            if (cur[1] + 1 < n) {
                heap.offer(new int[]{cur[0], cur[1] + 1});
            }
        }
        // get the index of the kth smallest element
        int[] res = heap.poll();
        return matrix[res[0]][res[1]];
    }
}

复杂度分析

• 时间复杂度：O(k*logn)
• 空间复杂度：O(n)

[snmyj]

class Solution {
public:
    int kthSmallest(vector<vector<int>>& matrix, int k) {
        vector<int> ret;
        for(int i=0;i<matrix.size();i++){
            for(int j=0;j<matrix[0].size();j++){
                ret.push_back(matrix[i][j]);
            }
        }
        sort(ret.begin(),ret.end());
        return ret[k-1];
    }
};

[FlipN9]

/**
    Method 2: Binary Search         0ms 100%
    
    矩阵内的元素是从左上到右下递增
    那么任意取一个 mid, mid 会将矩阵分割成两块, 按照分割线从左下往右上走一遍, 即可算出 < mid 的数有多少个
            如果数量 > k, mid 应该缩小
                    < k, mid 应该变大
    
    找到 count == k 的点
    
    TC: O(nlog(r - l))      二分查找的次数为 log(max - min), 每次操作的 findNotBiggerThanMid 为 O(N)
    SC: O(1)
*/
class Solution {
    int[][] matrix;
    int R, C;
    
    public int kthSmallest(int[][] matrix, int k) {
        this.matrix = matrix;
        this.R = matrix.length;
        this.C = matrix[0].length;
        
        int left = matrix[0][0], right = matrix[R - 1][C - 1];
        while (left < right) {
            // 每次循环都保证第K小的数在[left, right]之间，当left==right，第k小的数就是left
            int mid = left + (right - left) / 2;
            int count = findNotBiggerThanMid(mid);
            if (count < k)  // count 的数量比 k 少, 
                left = mid + 1;
            else 
                right = mid;
        }
        return left; 
    }
    
    private int findNotBiggerThanMid(int mid) {
        // 左下开始找, 先行, 碰到 > mid 的, row--, 然后在从当前 col 开始继续往右
        int i = R - 1, j = 0;
        int count = 0;
        while (i >= 0 && j < C) {
            if (matrix[i][j] <= mid) {
                // 第j列有i+1个元素<=mid
                count += i + 1;
                j++;
            } else {
                // 第j列目前的数大于mid，需要继续在当前列往上面的 row找
                i--;
            }
        }
        return count;
    }
}


/**
    Method 1: PQ k路归并 54ms
    
    TC: O(k * logN)
    SC: O(N)
*/
class Solution1 {
    public int kthSmallest(int[][] matrix, int k) {
        int N = matrix.length;
        // <val, row, col>
        PriorityQueue<int[]> pq = new PriorityQueue<>((a, b) -> (a[0] - b[0]));
        for (int i = 0; i < N; i++) 
            pq.add(new int[] {matrix[i][0], i, 0});
        for (int i = 0; i < k - 1; i++) {
            int[] cur = pq.poll();
            if (cur[2] != N - 1) { // 不是一行的最末尾
                pq.offer(new int[] {matrix[cur[1]][cur[2] + 1], cur[1], cur[2] + 1});
            }
        }
        return pq.poll()[0]; //第 k 个
    }
}

[Elsa-zhang]

import heapq
class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        lists = []
        for m in matrix:
            lists+=m
        heapq.heapify(lists)
        for i in range(k-1):
            heapq.heappop(lists)
        return lists[0]

[klspta]

class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        int n = matrix.length;
        int left = matrix[0][0];
        int right = matrix[n - 1][n - 1];
        while (left < right) {
            int mid = left + ((right - left) >> 1);
            if (check(matrix, mid, k, n)) {
                right = mid;
            } else {
                left = mid + 1;
            }
        }
        return left;
    }

    public boolean check(int[][] matrix, int mid, int k, int n) {
        int i = n - 1;
        int j = 0;
        int num = 0;
        while (i >= 0 && j < n) {
            if (matrix[i][j] <= mid) {
                num += i + 1;
                j++;
            } else {
                i--;
            }
        }
        return num >= k;
    }
}

[bookyue]

code

1. binary search
2. minimum Heap

    /**
    * TC: O(nlog(r - l)
    * SC: O(1)
    */
    public int kthSmallest(int[][] matrix, int k) {
        int n = matrix.length;
        if (n * n == k) return matrix[n - 1][n - 1];

        int left = matrix[0][0];
        int right = matrix[n - 1][n - 1];

        while (left < right) {
            int mid = (left + right) >> 1; // for negative numbers, the results of division and signed right shift aren't the same
            if (enough(matrix, k, mid))
                right = mid;
            else 
                left = mid + 1;
        }

        return left;
    }

    private boolean enough(int[][] matrix, int k, int x) {
        int count = 0;
        int n = matrix.length;
        int j = n - 1;

        for (int[] row : matrix) {
            while (j >= 0 && row[j] > x) j--;
            count += j + 1;
        }

        return count >= k;
    }

    /**
    * TC: O(klogn)
    * OC: O(n)
    */
    public int kthSmallest(int[][] matrix, int k) {
        int n = matrix.length;
        if (n * n == k) return matrix[n - 1][n - 1];

        PriorityQueue<int[]> minHeap = new PriorityQueue<>(Comparator.comparingInt(a -> a[2]));
        for (int i = 0; i < Math.min(n, k); i++)
            minHeap.offer(new int[]{i, 0, matrix[i][0]});

        for (int i = 1; i < k; i++) {
            var p = minHeap.poll();
            
            if (p[1] < n - 1)
                minHeap.offer(new int[]{p[0], p[1] + 1, matrix[p[0]][p[1] + 1]});
        }

        return minHeap.poll()[2];
    }

[tiandao043]

思路

一开始想的直接排序，然后行有序想到归并。
看了题解，列也有序，二分，但不是用索引，找值。

代码

class Solution {
public:
    bool check(vector<vector<int>>& matrix, int mid, int k, int n) {
        int i = n - 1;
        int j = 0;
        int num = 0;
        while (i >= 0 && j < n) {
            if (matrix[i][j] <= mid) {
                num += i + 1;
                j++;
            } else {
                i--;
            }
        }
        return num >= k;
    }
    
    int kthSmallest(vector<vector<int>>& matrix, int k) {
        int n=matrix.size();
        int l=matrix[0][0];
        int r=matrix[n-1][n-1];
        while(l<r){
            int mid=l+((r-l)>>1);
            if(check(matrix,mid,k,n)){
                r=mid;
            }else{
                l=mid+1;
            }
        }
        return l;

    }
};

[Jetery]

class Solution {
public:
    int kthSmallest(vector<vector<int>>& matrix, int k) {
        int l = matrix[0][0], r = matrix[matrix.size() - 1][matrix.size() - 1];
        while(l <= r){
            int m = (l + r) >> 1, count = 0, index = 0;
            for(int i = matrix.size() - 1; i >= 0; i--){
                while(index < matrix.size() && matrix[i][index] <= m){
                    index++;
                }
                count += index;
            }
            if(count < k)l = m + 1;
            else r = m - 1;
        }
        return l;
    }
};

[mayloveless]

思路

二分查找，检查每次找到的值是不是第k个

代码

/**
 * @param {number[][]} matrix
 * @param {number} k
 * @return {number}
 */
var kthSmallest = function(matrix, k) {
    const n = matrix.length;
    const check = function(mid) {
        let i = n-1;
        let j = 0;
        let res = 0;
        while (i >= 0 && j < n){
            if (mid >= matrix[i][j]){
                res += (i + 1)
                j += 1
            } else{
                i -= 1
            }   
        }
        return res >= k
    };
    let left = matrix[0][0]; // 左指针
    let right  = matrix[n-1][n-1]; // 右指针
    while (left < right) { // 相等即退出
        let mid = (left + right) >>> 1; //无符号右移，避免溢出。此处取左中位
        if( !check(mid) ){  // target根据具体场景需求
            left = mid + 1; //左边舍去，如果取左中位，则需要左边指针收缩，避免死循环。
        }else{
            right = mid;//右边舍去
        }
    }
    return left;// 左右相等都可以。
};

复杂度

时间O(nlogn)
空间O（1）

[RestlessBreeze]

code

class Solution {
public:
    bool check(vector<vector<int>>& matrix, int mid, int k, int n) {
        int i = n - 1;
        int j = 0;
        int num = 0;
        while (i >= 0 && j < n) {
            if (matrix[i][j] <= mid) {
                num += i + 1;
                j++;
            } else {
                i--;
            }
        }
        return num >= k;
    }

    int kthSmallest(vector<vector<int>>& matrix, int k) {
        int n = matrix.size();
        int left = matrix[0][0];
        int right = matrix[n - 1][n - 1];
        while (left < right) {
            int mid = left + ((right - left) >> 1);
            if (check(matrix, mid, k, n)) {
                right = mid;
            } else {
                left = mid + 1;
            }
        }
        return left;
    }
};