Leetcode 378. Kth Smallest Element in a Sorted Matrix

Posted 2020-08-24 lettuan

Given a n x n matrix where each of the rows and columns are sorted in ascending order, find the kth smallest element in the matrix.

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

Example:

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

return 13.

本题的解法一和Find K pairs with smallest sums完全类似。在那一题中，每当(A1,Bk)小于堆顶元素时，我们就将所有的(A1,Bk),...,(An,Bk)推到堆里面。然后idx2 += 1。这里我们做类似的事情，每当matrix[row][0]小于堆顶元素时，我们就将整个一行都推到heap里面。还是要注意heap里需要预存一个MaxInt，否则可能出现pop空堆的情况。

 1 class Solution(object):
 2     def kthSmallest(self, matrix, k):
 3         """
 4         :type matrix: List[List[int]]
 5         :type k: int
 6         :rtype: int
 7         """
 8         ans = []
 9         heap = [(0x7FFFFFFF, None, None)]
10         size1, size2 = len(matrix), len(matrix[0])
11         row = 0 
12         while len(ans) < k:
13             if row < size1:
14                 if matrix[row][0] < heap[0]:
15                     for col in range(size2):
16                         heapq.heappush(heap, matrix[row][col])
17                     row += 1
18             val = heapq.heappop(heap)
19             ans.append(val)
20         return ans[k-1]

 

其实由于本题的矩阵的特殊结构，最小的元素肯定是matrix[0][0]，我们只需要维护一个Heap，每次pop出堆顶元素的时候，把大于他的最小的两个元素（[row, col+1] 和 [row+1, col]）推入堆。当heappop执行第k次时，pop出来的那个值就是矩阵中第k小的元素。这里需要维护一个visited矩阵来记录一个节点是不是已经推入heap中。否则的话会出现一个点被多次推入Heap的情况。

 1 class Solution(object):
 2     def kthSmallest(self, matrix, k):
 3         """
 4         :type matrix: List[List[int]]
 5         :type k: int
 6         :rtype: int
 7         """
 8         m, n = len(matrix), len(matrix[0])
 9         heap = [(matrix[0][0], 0, 0)]
10         visited = [[False]*n for _ in range(m)]
11 
12         for i in range(k):
13             val, row, col = heap[0]
14             visited[row][col] = True
15             if col < n-1 and visited[row][col+1] == False:
16                 visited[row][col+1] = True
17                 heapq.heappush(heap, (matrix[row][col+1], row, col+1))
18             if row < m-1 and visited[row+1][col] == False:
19                 visited[row+1][col] = True
20                 heapq.heappush(heap, (matrix[row+1][col], row+1, col))
21             ans, row, col = heapq.heappop(heap)
22         
23         return ans

 解法三，二分法查找 binary search

左上角的元素low = matrix[0][0]，和右下角的元素high = matrix[-1][-1]显然是是最小和最大的。我们可以查找有多少个元素小于等于 mid = (low + high)//2。

考虑如下的例子： low = 1, high = 22, mid = 11

从左下角开始寻找，如果matrix[i][j] > mid的话，那么它右边的元素肯定也都大于mid。我们就查看上一行(i--)，如果matix[i][j] <= mid的话，那么它本身和它上面的元素都<=mid。所以cnt += (i+1), 并且j ++。边界条件就是(i>=0, j<=n-1)。图中的case，小于等于mid的元素个数就是4+4+2+2+2= 14个。

为了方便我们把第k个元素写成L。如果14<k, 那么说明 mid=11 比 L 要严格小，可以设置 low = mid+1。反之说明 mid>= L，可以设置high = mid。

 

 1     def kthSmallest(self, matrix, k):
 2         """
 3         :type matrix: List[List[int]]
 4         :type k: int
 5         :rtype: int
 6         """
 7         n = len(matrix)
 8         low, high = matrix[0][0], matrix[-1][-1]
 9         while low < high:
10             mid = (low + high)>>1
11             temp = self.search_lower_than_mid(matrix, n, mid)
12             if temp < k:
13                 low = mid + 1
14             else:
15                 high = mid
16         return low
17         
18     def search_lower_than_mid(self, matrix, n, x):
19         i, j = n - 1, 0
20         cnt = 0
21         while i >= 0 and j < n:
22             if matrix[i][j] > x:
23                 i -= 1
24             else:
25                 j += 1
26                 cnt += i+1
27         return cnt

代码中留意L9, 这里不能用 low <= high。这样的话当最后high = low的时候，程序会死循环。L16 return low或者return high都可以。