題解

MinHeap (Priority Queue)

class MinHeap:
    def __init__(self):
        self.heap = []

    @property
    def length(self):
        return len(self.heap)

    def insert(self,val: int):
        self.heap.append(val)
        self.__swin(self.length - 1)
    
    def pop(self):
        min_value = self.heap[0]
        self.__swap(0,self.length - 1)
        self.heap.pop()
        self.__sink(0)
        return min_value
    
    def __swin(self,k: int):
        while k > 0 and self.heap[self.__parent(k)] > self.heap[k]:
            self.__swap(self.__parent(k),k)
            k = self.__parent(k)
    
    def __sink(self,k: int):
        while k * 2 + 1 < self.length:
            left,right = k * 2 + 1, k * 2 + 2
            target = left
            if right < self.length and self.heap[left] > self.heap[right]:
                target = right
            if self.heap[k] <= self.heap[target]:
                break
            self.__swap(k,target)
            k = target
    
    def __parent(self,k):
        return (k - 1) // 2
    
    def __swap(self, m: int, n: int):
        self.heap[m],self.heap[n] = self.heap[n], self.heap[m]

class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        n = len(matrix)
        min_heap = MinHeap()

        for row in matrix:
            for col in row:
                min_heap.insert(col)
        
        output = 0
        for i in range(k):
            output = min_heap.pop()
        
        return output

Binary Search (Martix)

因為 Matrix 的 X,Y 軸已經排序好，可以利用這個特性去做 Binary search

reference: Selection in X+Y and matrix with sorted rows and columns

class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        m, n = len(matrix), len(matrix[0])

        def check(mid):
            i, j = m - 1, 0
            nums = 0
            while i >= 0 and j < n:
                if matrix[i][j] <= mid: # 如果小於等於 mid 就代表是 kth 小的元素
                    nums += i + 1
                    j += 1
                else: # 不是就往上
                    i -= 1
            return nums >= k # 查看是否大於 kth

        left, right = matrix[0][0], matrix[-1][-1] # 因為要回傳 matrix 裡面的數字，所以用 matrix 裡面的數字直接做 left right
        while left < right:
            mid = left + (right - left) // 2
            if check(mid):
                right = mid
            else:
                left = mid + 1
        return left