378. Kth Smallest Element in a Sorted Matrix

Given an n x n matrix where each of the rows and columns is sorted in ascending order, return the kth smallest element in the matrix.

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

You must find a solution with a memory complexity better than O(n2).

Example 1:

Input: matrix = [[1,5,9],[10,11,13],[12,13,15]], k = 8
Output: 13
Explanation: The elements in the matrix are [1,5,9,10,11,12,13,13,15], and the 8th smallest number is 13

Example 2:

Input: matrix = [[-5]], k = 1
Output: -5

思路Coding化:

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

        '''
        [[1,5,9],[10,11,13],[12,13,15]]

        首先先把第一個數放到minheap
        # 1 (0,0)
        # 10 (1,0)
        # 12 (2,0)

        接著pop出最小的數後，把那個數的下一個數push進來
        res = [1]
        # 5 (0,1)
        # 10 (1,0)
        # 12 (2,0)

        res = [1,5]
        # 9 (0,2)
        # 10 (1,0)
        # 12 (2,0)
        依此類推
        '''
        minheap = []

        for r in range(min(k,len(matrix))):
            minheap.append([matrix[r][0],r,0])
        
        while k :
            res, r, c = heapq.heappop(minheap)
            #注意設定範圍 c+1時才不會out of range
            if c < len(matrix) -1:
                heapq.heappush(minheap, [matrix[r][c+1], r, c+1] )
            k-=1
        return res
#T: O(k log k)
#S: O(k)

Time Complexity:O(k log k)

Space Complexity:O(k)