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

파이썬으로 정렬된 리스트를 유지하며 요소를 삽입할 인덱스 찾는 방법

오름차순으로 정렬된 숫자 리스트 nums와 하나의 숫자 target이 주어졌다고 가정해 봅시다. 이때 nums의 정렬 상태가 유지되도록 target을 삽입해야 할 인덱스를 찾아야 합니다. 만약 target이 이미 리스트에 존재한다면, 삽입 가능한 인덱스 중 가장 뒤쪽(가장 큰) 인덱스를 반환해야 합니다.

추가 조건은 두 가지입니다. 라이브러리 함수를 사용하지 않고 직접 구현해야 하며, 시간 복잡도는 O(log n)을 만족해야 합니다. 이 조건을 모두 충족하는 방법은 바로 이진 탐색(Binary Search)입니다.

문제 예시

예를 들어 입력이 다음과 같다고 해 보겠습니다.

  • nums = [1,5,6,6,8,9]
  • target = 6

이 경우 출력은 4가 됩니다. 6이 이미 리스트에 존재하기 때문에, 삽입 가능한 가장 큰 인덱스인 4를 반환하면 최종 배열은 [1,5,6,6,6,8,9]처럼 여전히 정렬된 상태를 유지합니다.

풀이 접근 방식

이진 탐색을 활용하면 다음 단계로 문제를 해결할 수 있습니다.

  • left를 0으로, right를 리스트 길이에서 1을 뺀 값으로 초기화합니다.
  • ans를 0으로 초기화합니다.
  • left <= right인 동안 반복합니다.
    • mid를 (left + right) / 2의 내림값으로 설정합니다.
    • target >= nums[mid]라면, ans를 mid + 1로 갱신하고 left를 mid + 1로 이동합니다.
    • 그렇지 않다면, right를 mid - 1로 이동합니다.
  • 반복이 끝나면 ans를 반환합니다.

핵심 아이디어는 target이 중간값보다 크거나 같으면 삽입 위치는 반드시 오른쪽 영역에 있다는 점입니다. 이때 ans에 후보 인덱스를 저장해 두면, 동일한 값이 여러 개 있어도 자동으로 가장 뒤쪽 인덱스가 기록됩니다.

구현 코드

아래 파이썬 구현을 통해 더 잘 이해해 보겠습니다.

def solve(nums, target):
   left, right = 0, len(nums) - 1
   ans = 0
   while left <= right:
      mid = (left + right) // 2
      if target >= nums[mid]:
         ans = mid + 1
         left = mid + 1
      else:
         right = mid - 1
   return ans

nums = [1,5,6,6,8,9]
target = 6
print(solve(nums, target))

입력

[1,5,6,6,8,9], 6

출력

4

복잡도 분석

매 반복마다 탐색 범위가 절반으로 줄어들기 때문에 시간 복잡도는 O(log n)이며, 추가적인 자료 구조를 사용하지 않으므로 공간 복잡도는 O(1)입니다. 참고로 파이썬 표준 라이브러리의 bisect.bisect_right 함수가 동일한 동작을 수행하지만, 이번 문제에서는 직접 구현이 요구되므로 위와 같이 작성했습니다.