Python으로 정렬된 2차원 행렬에서 목표 값 찾기: 효율적인 탐색 알고리즘 구현

2차원 행렬이 있고, 각 행과 각 열이 오름차순(비내림차순)으로 정렬되어 있다고 가정해 보겠습니다. 이때 주어진 목표 값(target)이 이 행렬 안에 존재하는지 확인하는 프로그램을 작성해야 합니다.

예를 들어 입력이 다음과 같다고 가정해 봅시다.

그리고 target = 31이라면, 출력 결과는 True가 됩니다.

문제 해결 접근 방법

이 문제는 행렬의 오른쪽 위 모서리에서 탐색을 시작하는 '계단식 탐색(staircase search)' 기법으로 효율적으로 해결할 수 있습니다. 행과 열이 모두 정렬되어 있기 때문에, 현재 위치의 값이 목표 값보다 크면 열 인덱스를 줄이고, 그렇지 않으면 다음 행으로 넘어가는 방식으로 탐색 범위를 좁혀 나갈 수 있습니다.

구체적인 알고리즘 단계는 다음과 같습니다.

• col := 행렬의 열 개수 - 1 (마지막 열 인덱스)
• i를 0부터 행렬의 행 개수까지 반복합니다.
  • matrix[i][col] > target이고 col >= 0인 동안 다음을 반복합니다.
    • col := col - 1
  • 만약 matrix[i][col] == target이라면,
    • True를 반환합니다.
• 모든 반복이 끝나면 False를 반환합니다.

이 알고리즘의 시간 복잡도는 O(m + n)입니다. 여기서 m은 행의 개수, n은 열의 개수로, 전체 행렬을 하나하나 순회하는 O(m × n) 방식보다 훨씬 효율적입니다.

더 잘 이해할 수 있도록 다음 구현 예제를 살펴보겠습니다.

예제 코드

class Solution:
    def solve(self, matrix, target):
        col = len(matrix[0]) - 1
        for i in range(len(matrix)):
            while matrix[i][col] > target and col >= 0:
                col = col - 1
            if matrix[i][col] == target:
                return True
        return False

ob = Solution()
matrix = [[2, 4, 30], [3, 4, 31], [6, 6, 32]]
target = 31
print(ob.solve(matrix, target))

입력

matrix = [
[2, 4, 30],
[3, 4, 31],
[6, 6, 32]]
target = 31

출력

True