【Day 89 】2024-07-05 - 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 。

[CathyShang]

class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        n = len(matrix)
        h = [(matrix[i][0], i, 0) for i in range(n)]
        heapify(h)
        while k:
            num,x,y = heappop(h)
            if y!=n-1:
                heappush(h, (matrix[x][y+1], x, y+1))
            k -= 1
        ans = num
        return ans 

• time O(k)
• space O(1)

[Dtjk]

class Solution { public: int kthSmallest(vector<vector>& 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; } };

[zhiyuanpeng]

class Solution:

    def countLessEqual(self, matrix, mid, smaller, larger):

        count, n = 0, len(matrix)
        row, col = n - 1, 0

        while row >= 0 and col < n:
            if matrix[row][col] > mid:

                # As matrix[row][col] is bigger than the mid, let's keep track of the
                # smallest number greater than the mid
                larger = min(larger, matrix[row][col])
                row -= 1

            else:

                # As matrix[row][col] is less than or equal to the mid, let's keep track of the
                # biggest number less than or equal to the mid

                smaller = max(smaller, matrix[row][col])
                count += row + 1
                col += 1

        return count, smaller, larger

    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:

        n = len(matrix)
        start, end = matrix[0][0], matrix[n - 1][n - 1]
        while start < end:
            mid = start + (end - start) / 2
            smaller, larger = (matrix[0][0], matrix[n - 1][n - 1])

            count, smaller, larger = self.countLessEqual(matrix, mid, smaller, larger)

            if count == k:
                return smaller
            if count < k:
                start = larger  # search higher
            else:
                end = smaller  # search lower

        return start