Computer >> 컴퓨터 >  >> 프로그래밍 >> Python

파이썬으로 정렬된 행렬에서 K번째로 작은 요소 찾는 방법

n×n 크기의 행렬이 있고, 각 행과 각 열이 모두 오름차순으로 정렬되어 있다고 가정해 봅시다. 이때 이 행렬에서 k번째로 작은 요소를 찾아야 합니다. 주의할 점은 여기서 말하는 k번째가 '정렬된 순서' 기준이라는 것이며, 중복을 제거한 후의 k번째 고유 값이 아니라는 점입니다.

예를 들어 입력이 [[1,5,9],[10,11,13],[12,13,15]]이고 k = 8이라면, 전체 요소를 정렬했을 때 8번째 위치에 있는 값인 13이 출력됩니다.

문제 해결 접근 방식

이 문제는 이진 탐색(Binary Search)행렬 탐색을 결합하여 효율적으로 해결할 수 있습니다. 알고리즘의 단계는 다음과 같습니다.

1. checkVal() 메서드 정의

  • checkVal()이라는 메서드를 정의하고, 인자로 행렬(matrix)과 기준값(value)을 받습니다.
  • i := 0, j := matrix[0]의 길이 − 1, counter := 0으로 초기화합니다.

2. 행렬 탐색으로 개수 세기

  • i가 행렬의 길이보다 작고 j가 0 이상인 동안 반복합니다.
  • 만약 matrix[i][j] > value라면 j를 1 감소시킵니다.
  • 그렇지 않다면 counter에 j + 1을 더하고 i를 1 증가시킵니다.
  • 반복이 끝나면 counter를 반환합니다. 이 값은 기준값보다 작거나 같은 요소의 개수입니다.

3. 메인 메서드에서 이진 탐색 수행

  • n := 행렬의 행 개수, high := 우하단(오른쪽 아래) 요소, low := 좌상단(왼쪽 위) 요소로 설정합니다.
  • low ≤ high인 동안 반복합니다.
  • mid = low + (high − low) / 2로 중간값을 계산합니다.
  • count := checkVal(matrix, mid)로 mid 이하의 요소 개수를 구합니다.
  • count < k라면 low = mid + 1로 갱신하고, 그렇지 않으면 high = mid − 1로 갱신합니다.
  • 반복이 종료되면 low를 반환합니다. 이것이 바로 k번째로 작은 요소입니다.

아래 예제 코드를 통해 더 자세히 이해해 보겠습니다.

구현 예제 코드

class Solution(object):
    def kthSmallest(self, matrix, k):
        """
        :type matrix: List[List[int]]
        :type k: int
        :rtype: int
        """
        n = len(matrix)
        high = matrix[n-1][n-1]
        low = matrix[0][0]
        while low<=high:
            mid = low + (high - low) /2
            count = self.check_value(matrix,mid)
            if count< k:
                low = mid+1
            else :
                high = mid-1
        return int(low)
    def check_value(self, matrix, value):
        i = 0
        j = len(matrix[0])-1
        counter = 0
        while(i<len(matrix) and j >=0):
            if matrix[i][j] > value:
                j-=1
            else:
                counter+=j+1
                i+=1
        return counter

matrix = [[1,5,9],[10,11,13],[12,13,15]]
ob = Solution()
print(ob.kthSmallest(matrix, 8))

입력

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

출력

13

알고리즘 동작 원리 정리

이 알고리즘의 핵심은 두 가지입니다. 첫째, checkVal() 메서드는 행렬의 오른쪽 위 모서리에서 시작해 왼쪽 아래 방향으로 탐색하면서 O(n) 시간 안에 특정 값 이하의 요소 개수를 셀 수 있습니다. 둘째, 메인 메서드는 행렬의 최솟값(좌상단)과 최댓값(우하단) 사이에서 이진 탐색을 수행하며, 'k개 이상의 요소가 mid 이하인 가장 작은 값'을 찾아냅니다.

전체 시간 복잡도는 O(n log(max−min))로, 행렬의 모든 요소를 정렬하는 O(n² log n²) 방식보다 훨씬 효율적입니다. 공간 복잡도 역시 O(1)로 추가 메모리 사용 없이 문제를 해결할 수 있다는 장점이 있습니다.