go

Author: robot

problems/378.kth-smallest-element-in-a-sorted-matrix.md

@@ -72,6 +72,8 @@ k = 8,
 
 - 最后直接返回 start, end, 或者 mid 都可以，因此三者最终会收敛到矩阵中的一个元素，这个元素也正是我们要找的元素。
 
+关于如何计算 count，我们可以从左下或者右上角开始，每次移动一个单位，直到找到一个值大于等于中间值，然后计算出 count，具体见下方代码。
+
 整个计算过程是这样的：
 
 ![378.kth-smallest-element-in-a-sorted-matrix-4](https://tva1.sinaimg.cn/large/007S8ZIlly1ghltyfm2okj30je0fq0uf.jpg)
@@ -92,6 +94,10 @@ k = 8,
 
 ## 代码
 
+代码支持：JS，Python3，CPP
+
+JS：
+
 ```js
 /*
  * @lc app=leetcode id=378 lang=javascript
@@ -146,6 +152,77 @@ var kthSmallest = function (matrix, k) {
 };
 ```
 
+Python3：
+
+```python
+class Solution:
+    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
+        n = len(matrix)
+
+        def check(mid):
+            row, col = n - 1, 0
+            num = 0
+            while row >= 0 and col < n:
+                # 增加 col
+                if matrix[row][col] <= mid:
+                    num += row + 1
+                    col += 1
+                # 减少 row
+                else:
+                    row -= 1
+            return num >= k
+
+        left, right = matrix[0][0], matrix[-1][-1]
+        while left <= right:
+            mid = (left + right) // 2
+            if check(mid):
+                right = mid - 1
+            else:
+                left = mid + 1
+
+        return left
+
+```
+
+CPP Code:
+
+```cpp
+class Solution {
+public:
+    bool check(vector<vector<int>>& matrix, int mid, int k, int n) {
+        int row = n - 1;
+        int col = 0;
+        int num = 0;
+        while (row >= 0 && col < n) {
+            if (matrix[row][col] <= mid) {
+                num += i + 1;
+                col++;
+            } else {
+                row--;
+            }
+        }
+        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 - 1;
+            } else {
+                left = mid + 1;
+            }
+        }
+        return left;
+    }
+};
+
+
+```
+
 **复杂度分析**
 
 - 时间复杂度：二分查找进行次数为 $O(log(r-l))$，每次操作时间复杂度为 O(n)，因此总的时间复杂度为 $O(nlog(r-l))$。