LeetCode의 Find First and Last Position of Element in Sorted Array 문제다.
정렬된 오름차순 배열에서 특정 값이 어디부터 어디까지 있는지 탐색하는 문제다. 예를 들어 [1, 2, 3, 3, 3, 4, 5]가 주어진다면 3을 탐색한 결과가 [2, 4]처럼 반환되어야 한다. 왜냐면 인덱스 2부터 인덱스 4까지가 3이기 때문이다.
별다른 생각 없이 그냥 단순하게 파이썬의 index 메서드를 이용하여 푸는 방식이다. 하지만 반복되는 리스트 분할 및 합병, 효과적이지 못한 index 메서드 활용 때문에 1초가 넘게 걸리는 매우 비효율적인 결과를 얻을 수 있었다.
class Solution:
def searchRange(self, nums: List[int], target: int) -> List[int]:
if target not in nums:
return [-1, -1]
firstIndex = nums.index(target)
findIndex = firstIndex
lastIndex = -1
while findIndex < len(nums) and target in nums[findIndex:]:
lastIndex = nums[findIndex:].index(target) + findIndex
findIndex = lastIndex+1
return [firstIndex, lastIndex]
분명히 올바른 풀이는 아니었기 때문에 다시 곰곰히 생각해보았고 정렬된 리스트에서 탐색하는데 매우 효과적인 이진 탐색을 떠올릴 수 있었다.
정렬된 배열이기 때문에 그냥 탐색해서 시작과 끝을 찾으면 될 것 같지만 문제에서는 O(logn) 시간 복잡도를 요구하고 있다. 즉 이진 탐색으로 해결하라는 힌트를 준 것이다.
이진 탐색은 정렬된 리스트에서 빠르게 값을 찾을 수 있는 탐색 방법 중 하나로 위키피디아의 이진 탐색 알고리즘에서 다음과 같이 정의되어 있다.
알고리즘만 보면 언제 탐색이 종료되는지 헷갈릴 수 있는데 M에서 원소를 찾지 못했을 때 L과 R을 하나씩 갱신하면서 범위를 좁혀가는 것을 볼 수 있다. 예를 들어 [1, 2, 3, 4, 5]에서 4를 찾는 경우를 생각해보자. L과 R은 각각 (0, 4)에서 시작한다.
첫 번째 탐색에서 M 번째 요소(3) 뒤에 4가 있기 때문에 (M+1, R) = ((L + R) / 2 + 1, R) = ((0 + 4) / 2 + 1, 4) = (3, 4)로 이동하게 된다.
다음 탐색에서는 M번째 요소(3) 뒤에 4가 있기 때문에 (M+1, R) = (FLOOR((3 + 4) / 2) + 1, 4) = (3+1, 4) = (4, 4)로 이동한다.
M 번째 요소(4)는 찾고자 하는 요소를 정확히 가리키기 때문에 탐색을 종료한다.
위처럼 탐색 범위를 절반으로 줄여가면서 탐색하는 것을 볼 수 있다. 이를 기반으로 작성한 코드는 다음과 같다.
class Solution:
def searchRange(self, nums: List[int], target: int) -> List[int]:
# 기본적인 예외처리
if len(nums) == 0:
return [-1, -1]
if len(nums) == 1:
return [0, 0] if nums[0] == target else [-1, -1]
# 이진 탐색.
left = -1
low, high = 0, len(nums)-1
mid = 0
while low <= high:
mid = int((low + high) / 2)
if target < nums[mid]:
high = mid - 1
elif target > nums[mid]:
low = mid + 1
else:
break
# 탐색 대상이 존재하지 않는다면 [-1, -1] 반환.
if nums[mid] == target:
left = mid
else:
return [-1, -1]
right = mid
# 왼쪽 원소까지 탐색.
while left > 0 and nums[left-1] == target:
left = left - 1
# 오른쪽 원소까지 탐색.
while right < len(nums)-1 and nums[right+1] == target:
right = right + 1
return [left, right]
이진 탐색 알고리즘과는 관련 없는 반복문도 두 개 추가되었는데 이는 문제에서 단순히 원소가 해당 리스트에 존재하는지 뿐 아니라 어디부터 어디까지 존재하는지 탐색해야 하기 때문에 자신과 다른 원소를 발견할 때까지 탐색해야 한다. 이 부분이 O(n)에 걸리지 않을까 걱정되긴 하지만 지식이 부족하기 때문에 더 나은 개선점은 생각해내지 못했다.
이진 탐색을 구현하는 경험은 오랜만이라 까다로웠는데 특히 탐색의 종료 조건이나 L, R 인덱스를 M으로 할지 M+1로 할지 결정하기 못해서 코드를 계속 잘못 작성했었다. 예전에 구현했던 경험이 부분부분 남아서 혼란스러웠는데 위키피디아에서 잘 정리된 알고리즘을 보고 명확하게 구현하는게 많은 도움이 되는것 같다.