[알고리즘, Python] 회전 정렬 배열 탐색 (Search in Rotated Sorted Array)

문제는 leecode에서 확인할 수 있다.

https://leetcode.com/problems/search-in-rotated-sorted-array/

Search in Rotated Sorted Array - LeetCode

---

이 문제를 풀기 전 기존에 정렬 배열에서 어떻게 탐색하는지 생각해보자.

 

탐색이므로 직관적으로 binary search를 사용해야겠다는 생각이 든다.

 

예를 들어 [1, 2, 4, 5, 7, 8, 9] 라는 배열이 있다고 가정하고 8의 index를 찾는다고 가정하면 코드는 아래와 같다.

def search(nums, target):
    low = 0
    high = len(nums)-1
    
    while low <= high:
    	mid = int((low+high)/2)
        if nums[mid] == target:
            return mid
        elif nums[mid] > target:
            high = mid-1
       	else:
            low = mid+1
    return -1

하지만 이 leecode 문제의 경우 알 수 없는 피봇으로 인해 회전되어 있는 정렬된 배열이 주어진다. (예: [4, 1, 2, 3])

 

따라서 target이 가운데를 기점으로 왼쪽에 있을지 오른쪽에 있을지를 잘 판단해야한다.

def search(self, nums: List[int], target: int) -> int:
    low = 0
    high = len(nums) - 1

    while low <= high:
        mid = int((low + high) / 2)
        if nums[mid] == target:
            return mid

        if nums[low] <= nums[mid]:                             #1
            if target >= nums[low] and target < nums[mid]:
                high = mid - 1
            else:
                low = mid + 1
        else:                                                  #2
            if target > nums[mid] and target <= nums[high]:
                low = mid + 1
            else:
                high = mid - 1
    return -1

우선 내가 탐색하려는 구간에서 왼쪽(low~mid)을 살펴본다.

 

(#1)

만약 [1, 2, 3, 4]와 같이 low 인덱스의 원소가 mid 인덱스의 원소보다 작은 경우 그냥 정렬된 배열임을 알 수 있다.

 

따라서 이 경우 타겟이 low와 mid 사이에 있으면 왼쪽을 탐색하고 그게 아니라면 오른쪽을 탐색하면 된다.

예1)

nums: [1, 2, 3, 4, 5] / target: 4 인 경우 4는 1~3 사이에 없으므로 오른쪽 탐색

예2)

nums: [2, 3, 4, 5, 1] / target: 1인 경우 1이 2~4 사이에 없으므로 오른쪽 탐색

 

(#2)

만약 low 인덱스 원소가 mid 인덱스 원소보다 큰 경우 low~mid에 회전 구간이 포함되어 있음을 알 수 있다. (예: [8, 1, 2, 3, 5])

 

이 경우에는 mid~high 가 정렬된 구간이므로 target이 그 사이에 위치한다면 오른쪽을 탐색하고 아니라면 왼쪽을 탐색한다.

---

문제에서 요구한대로 시간복잡도 O(log n)로 문제를 해결했지만 성능이 좋지는 않았다.

더 나은 방법이 있나 알아봐야겠다.