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

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

題目來源：力扣（LeetCode）https://leetcode-cn.com/problems/kth-smallest-element-in-a-sorted-matrixpython

題目

給定一個 n x n 矩陣，其中每行和每列元素均按升序排序，找到矩陣中第 k 小的元素。請注意，它是排序後的第 k 小元素，而不是第 k 個不一樣的元素。數組

示例：bash

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

返回 13。

提示：微信

你能夠假設 k 的值永遠是有效的，1 ≤ k ≤ n2 。

解題思路

先審題，題目中給出的是二維數組，而問題須要求得的是二維數組中第 k 小的元素。在這裏，最直接的方法就是將二維數組轉換爲一維數組，對於轉換後的一維數組進行升序排序，那麼此時的一維數組中的第 k 個元素就是所求的結果。代碼大體以下：app

# python
def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
    res = []
    for row in matrix:
        for ele in row:
            res.append(ele)
    res.sort()
    return res[k-1]

在這裏，代碼並無使用矩陣的特性。而是轉換爲一維數組進行求解。這個解法的時間複雜度爲 O(n^2logn)，便是對 n x n 個數進行排序。code

由於上面的解法是將矩陣轉換爲一維數組，並非在矩陣的基礎上解決問題。下面使用二分查找，來嘗試以在矩陣的前提下，對問題進行分析解答。blog

思路：二分查找排序

考慮使用二分查找的緣由，由於題目中說明，矩陣中，每行每列元素都是升序排序的。element

也就是說在題目給定的矩陣當中，元素由左上到右下是遞增的，以示例爲基礎擴充展開分析：leetcode

示例：

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

將上面的示例稍微擴充以下（僅爲更方便展現二分查找方法的效果）：

matrix = [
   [ 1,  5,  9, 10, 11],
   [10, 11, 13, 14, 15],
   [12, 13, 15, 16, 17],
   [13, 14, 16, 17, 18],
   [14, 18, 22, 26, 30]
],
k = 8

將上面二維數組轉換爲下圖：

咱們假設左上角的元素爲 matrix[0][0]，右下角的元素爲 matrix[n-1][n-1]，根據題意以及上面的示圖可知。matrix[0][0] 是二維數組中的最小值，matrix[n-1][n-1] 是最大值。

咱們如今使用二分查找的方法來分析，設 matrix[0][0] 和 matrix[n-1][n-1] 分別爲 left 和 right。

如今咱們嘗試取 mid（left <= mid <= right），能夠發如今矩陣當中小於等於 mid 值的元素會分佈在矩陣的左上部分，而大於 mid 值的元素則分佈在矩陣的右下部分。例以下圖所示，此時取 mid 爲 15：

在這裏大於 mid 和小於等於 mid 的元素分爲兩部分，沿着紅色的分割線將二者分開。此時，咱們就能夠看到，大於 mid 和小於等於 mid 二者元素的個數分別有多少。

上面紅色分割線劃分的依據，在這裏進行分析：

由於每行每列的元素都是升序排列的，咱們前面的分析，元素須要與 mid 進行比較。例如，咱們如今要求分佈在左邊部分的元素，也就是元素值小於等於 mid 的部分。

此時咱們能夠考慮從左下角開始出發，往右上角去找。這樣可以快速收縮範圍。具體的流程以下：

以左下角爲起始點，往右上角開始擴散
若是當前位置的元素值小於等於 mid 值時，說明從當前位置開始往上的全部元素都小於等於 mid（由於每列升序排列），記錄當前元素個數 count，注意維護更新。
此時知足元素值小於等於 mid 值時，向右邊移動再次比較。不然，就向上移動，尋找稍小的元素進行比較，直至移動到邊界。

在這裏，當取 mid 值後進行劃分時，按照上面的方法，可以肯定小於或等於 k 的個數，那麼就會出現如下兩種狀況：

當 count 大於或等於 k 時，那麼能夠肯定要找的答案在 [left, mid] 這邊；
不然，答案在 [mid+1, right] 這個區間。

那麼根據這個思路，循環直至找到答案。

關於二分查找方法執行流程，可看以下簡略圖示（若不太理解，可手畫圖幫忙理解）：

具體的代碼實現以下。

代碼實現

class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        def search(mid, n):
            # 從左下角開始往右上角找
            i, j = n-1, 0
            # 統計小於或等於 mid 的元素個數
            count = 0
            while i >= 0 and j < n:
                if matrix[i][j] <= mid:
                    # 當此時元素小於等於 mid
                    # 那麼該元素往上的全部元素都會小於或等於 mid
                    count += (i + 1)
                    # 向右移動，繼續比較
                    j += 1
                else:
                    # 當此時元素大於 mid 時，說明這個元素不包含在內，往上移動
                    i -= 1
            # count 與 k 值比對，判斷所求元素落在哪邊
            return count >= k
        
        n = len(matrix)

        left, right = matrix[0][0], matrix[n-1][n-1]

        while left < right:
            mid = left + (right - left) // 2
            if search(mid, n):
                # 當 count 大於等於 k 時，代表答案落在 [left, mid]
                # 不然落在 [mid+1, right]
                right = mid
            else:
                left = mid + 1
        
        return left

實現結果

文章原創，以爲寫得還能夠，歡迎關注點贊。微信公衆號《書所集錄》同步更新，一樣歡迎關注。

LeetCode 378. 有序矩陣中第K小的元素 | Python

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

題目

解題思路

代碼實現

實現結果