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

파이썬에서 배열 요소가 인덱스와 동일한 가장 작은 인덱스 찾기 (이진 탐색)

서로 중복되지 않는 고유한 요소들로 구성되어 있고 오름차순으로 정렬된 리스트 nums가 있다고 가정해 보겠습니다. 이때 nums[i] = i를 만족하는 최소 인덱스 i를 찾아야 하며, 조건을 만족하는 인덱스가 존재하지 않으면 -1을 반환해야 합니다. 또한 이 문제는 O(log n) 시간 복잡도 내에 해결해야 한다는 제약이 있습니다.

예를 들어 입력이 nums = [-4, -1, 2, 3, 8]이라면 결과는 2입니다. nums[2] = 2nums[3] = 3이 모두 조건을 만족하지만, 그중 더 작은 값이 2이기 때문입니다.

접근 방법: 이진 탐색 활용하기

O(log n)이라는 시간 제약 때문에 처음부터 끝까지 확인하는 단순 선형 탐색은 사용할 수 없습니다. 대신 이진 탐색(Binary Search)을 활용하면 효율적으로 문제를 해결할 수 있습니다.

핵심 아이디어는 다음과 같습니다. 리스트가 오름차순으로 정렬되어 있고 모든 요소가 고유한 정수라면, f(i) = nums[i] − i 는 절대 감소하지 않는 성질을 가집니다. 따라서 f(i) = 0, 즉 nums[i] = i가 되는 지점을 이진 탐색으로 빠르게 좁혀 나갈 수 있습니다.

알고리즘 단계

  • ret := -1, lhs := 0, rhs := len(nums) - 1로 초기화합니다.

  • lhs <= rhs인 동안 다음 과정을 반복합니다.

    • mid := (lhs + rhs) // 2로 중간 인덱스를 계산합니다.

    • nums[mid] == mid라면 → ret := mid로 저장합니다. (조건을 만족하는 후보를 발견)

    • nums[mid] >= mid라면 → rhs := mid - 1로 왼쪽 절반을 탐색합니다. (더 작은 인덱스 후보를 찾기 위해)

    • 그렇지 않다면 → lhs := mid + 1로 오른쪽 절반을 탐색합니다.

  • 반복이 종료되면 ret을 반환합니다. 조건을 만족하는 인덱스가 없었다면 초기값 -1이 그대로 반환됩니다.

구현 예제

다음 파이썬 코드를 통해 더 잘 이해해 보겠습니다.

def solve(nums):
   ret = -1
   lhs = 0
   rhs = len(nums) - 1
   while lhs <= rhs:
      mid = (lhs + rhs) // 2
      if nums[mid] == mid:
         ret = mid
      if nums[mid] >= mid:
         rhs = mid - 1
      else:
         lhs = mid + 1
   return ret

nums = [-4, -1, 2, 3, 8]
print(solve(nums))

입력

[-4, -1, 2, 3, 8]

출력

2

복잡도 분석

시간 복잡도: 매 반복마다 탐색 범위가 절반으로 줄어들므로 O(log n)입니다.
공간 복잡도: 추가적인 자료구조를 사용하지 않으므로 O(1)입니다.