[LeetCode] 378 | Kth Smallest Element in an Sorted Matrix | Python

---

문제 링크: https://leetcode.com/problems/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 $k^{th}$ smallest element in the matrix.

Note that it is the $k^{th}$ smallest element in the sorted order, not the $k^{th}$ distinct element.

You must find a solution with a memory complexity better than O($n^2$).

---

조건

• n == matrix.length == matrix[i].length
• 1 <= n <= 300
• $-10^9$ <= matrix[i][j] <= $10^9$
• All the rows and columns of matrix are guaranteed to be sorted in non-decreasing order.
• 1 <= k <= $n^2$

---

내 풀이

Sorting을 한 후에 k번째 원소를 찾으면 되지만, 뭔가 다른 방법으로 하고싶을 때, Binary Search를 사용하자.

left와 right를 각각 제일 작은 원소와 큰 원소로 설정하고, mid를 그 사이의 수로 설정한다.

mid에 대해서 모든 row에서 mid보다 작은 원소의 갯수를 세서 합친다 → 이 matrix에서 mid가 count번째 원소가 된다.

count가 k보다 작으면 조금 더 큰 수를 골라야하므로 left를 키우고, 반대의 상황이면 right를 줄인다.

---

코드

class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        N = len(matrix)
        left, right = matrix[0][0], matrix[N-1][N-1]
        
        while left < right:
            mid = left + (right-left)//2
            temp = N - 1
            count = 0
            
            for i in range(N):
                while temp >= 0 and matrix[i][temp] > mid:
                    temp -= 1
                
                count += (temp+1)
            
            if count < k:
                left = mid + 1
            else:
                right = mid
        
        return left