3 분 소요

Heap

문제

문제 링크

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

Constraints:

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

Follow up:

  • Could you solve the problem with a constant memory (i.e., O(1) memory complexity)?
  • Could you solve the problem in O(n) time complexity? The solution may be too advanced for an interview but you may find reading this paper fun.

스스로 고민

  • 가장 작다? 정렬이다? 그럼 힙이다!
  • 예전에는 그냥 삽질만 했었는데 확실히 사람은 머리에 뭘 넣어야 작동한다.
  • 참고로 min 또는 max를 활용해서 문제가 풀릴거 같으면 heap도 같이 떠올려주면 된다.

접근법

  • heap을 이용하지 않고 풀면 문제가 매우 쉽다.
class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        if len(matrix) == 0:
            return 0

        result = []
        for n in matrix:
            result += n
            
        result.sort()
        return result[k-1]
  • 항상 0으로 예외 케이스를 만드니 처음에 조건문을 실행해주고 나머지는 파이썬 특성을 이용해서 풀면 된다.

  • 하지만 최근에 힙을 잘 다루지 못했으니까 힙으로 문제를 풀어야한다.
    • 지금처럼 쑥쑥 잘 풀 수 있을 때까지 연습
class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        if len(matrix) == 0:
            return 0
        result = []
        
        for n in matrix:
            result += n

        heapify(result)
        for _ in range(k-1):
            heappop(result)

        return heappop(result)
  • 일단 힙은 최소 값만 보장될 뿐 두 번째도 연속으로 최소 값이 되지 는 않는다. 하지만 heappop으로 추출할 경우 알아서 정렬이 된다.
  • 또한 matrix가 2차원 행렬로 되어있어서 바로 힙에 넣으면 묶여서 움직이므로 반드시 풀어줘야 한다.

시간 복잡도와 공간 복잡도

  1. 시간 복잡도:
    • 먼저, 주어진 행렬(matrix)의 모든 요소를 1차원 리스트(result)에 병합한다. 이 작업에는 O(N*M) 시간이 걸린다. 여기서 N은 행의 수, M은 열의 수다.
    • 다음으로, 결과 리스트를 최소 힙(heap)으로 변환하는 데는 O(N*M) 시간이 소요된다.
    • 그런 다음, k-1 번의 반복문을 사용하여 k-1번째로 작은 요소까지 탐색합니다. 이 작업에는 O(k-1) 시간이 걸린다.
    • 마지막으로, 최소 힙에서 요소를 하나 더 팝하여 k번째로 작은 요소를 찾는다. 이 작업에는 O(log(NM)) 시간이 소요된다. 따라서 총 시간 복잡도는 O(N_M + k-1 + log(N_M))다. 일반적으로 N_M이 k보다 큰 경우, O(N_M)이 가장 큰 영향을 미치므로 시간 복잡도는 O(NM)다.
  2. 공간 복잡도:
    • 결과 리스트(result)는 원래 행렬(matrix)의 모든 요소를 포함하므로, 공간 복잡도는 O(N*M)다.
    • 최소 힙을 구성하는 데 필요한 공간도 O(N*M)다.
    • 그 외에 추가적인 공간을 사용하지 않으므로, 총 공간 복잡도는 O(N*M)다.

회고 과정

  • 언제 힙을 쓰면 적절한지는 아직까지 감이 잘 오지 않았다.
  • 해당 문제를 사람들이 다른 방법으로는 어떻게 풀었는지 찾아봤다.

다른 사람 코드를 보고 그 기준으로 어떤 부분을 개선 할 수 있는지 최종 회고

class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        if not matrix or not matrix[0]:
            return None

        rows, cols = len(matrix), len(matrix[0])
        left, right = matrix[0][0], matrix[rows - 1][cols - 1]

        while left < right:
            mid = (left+right)//2
            count = 0
            j = cols - 1

            for i in range(rows):
                while j >= 0 and matrix[i][j] > mid:
                    j -= 1
                count += j + 1

            if count < k:
                left = mid+1

            else:
                right = mid

        return left
  • 결과적으로 보면 맨 처음에 힙을 사용하지 않고 풀었던 방법과 성능차이는 크게 없다.
  • 이진 탐색으로 k 번째 작은 요소를 찾는다. 이 방식으로 시간 복잡도는 O(Nlog(M))으로 개선된다. 여기서 N은 행의 수, M은 열의 수다.
  • 행렬의 크기가 큰 경우에도 이진 탐색을 활용한 방식은 시간 복잡도가 효율적으로 관리되기 때문에 대부분의 경우에서 성능이 우수하다고 한다.

댓글남기기

참고