[LeetCode] Find First and Last Position of Element in Sorted Array

문제

https://leetcode.com/problems/find-first-and-last-position-of-element-in-sorted-array/

 

Given an array of integers nums sorted in ascending order, find the starting and ending position of a given target value. Your algorithm's runtime complexity must be in the order of \(O(\log n)\). If the target is not found in the array, return [-1, -1].

 

Example 1:

Input: nums = [5,7,7,8,8,10], target = 8

Output: [3,4]

 

Example 2:

Input: nums = [5,7,7,8,8,10], target = 6

Output: [-1,-1]

풀이

주어진 배열에서 특정값을 가지는 주소를 \(O(\log n)\)에 찾는다고 했으니 당연히 binary search를 제일 먼저 떠올려야 한다. Binary search를 사용하기로 작정하고 문제를 살펴보면, 두 번의 전형적인 binary search 알고리즘으로 답을 구할 수 있다는 것을 알게 될 것이다. Starting position을 찾기 위해 1번, ending position을 찾기 위해 1번이 필요하다. 시간 복잡도는 \(O(\log n)\)이고, 공간 복잡도는 \(O(1)\)이다.

 

좀 더 알고리즘을 개선할 여지가 없는지 고민해보자. 두 번의 binary search에서 중복된 작업이 수행된다. 수행 시작부터 주어진 target value를 처음 찾을 때까지 작업이 이에 해당한다. 따라서 이 작업을 한 번만 하도록 알고리즘을 수정하면 수행 시간이 줄어들 것이다.

 

알고리즘을 크게 보면 다음 세 부분으로 나눌 수 있다.

 

1. Binary search를 수행하며 target value를 가지는 index를 찾는다. 만약 없으면 [-1, -1]을 리턴한다.
2. 1번에서 남아있는 왼쪽 구역에 대해 binary search를 계속 수행하여 starting position을 찾는다.
3. 1번에서 남아있는 오른쪽 구역에 대해서도 binary search를 수행하여 ending position을 찾는다.

코드 (C++)

vector<int> searchRange(vector<int>& nums, int target) {
    vector<int> ret;

    int left = 0;
    int right = nums.size() - 1;
    int mid;
 
    /* Step 1 */
    while (left <= right) {
        mid = (left + right) / 2;
        if (nums[mid] == target) {
            break;
        } else if (nums[mid] < target) {
            left = mid + 1;
        } else {
            right = mid - 1;
        }
    }
    
    if (left > right) {
        ret.push_back(-1);
        ret.push_back(-1);
        return ret;
    }
    
    /* Step 2 */
    int start_left = left;
    int start_right = mid;
    while (start_left < start_right) {
        int start_mid = (start_left + start_right) / 2;
        if (nums[start_mid] < target) {
            start_left = start_mid + 1;
        } else {
            start_right = start_mid;
        }
    }

    /* Step 3 */
    int end_left = mid;
    int end_right = right;
    while (end_left < end_right) {
        int end_mid = (end_left - end_right) / 2 + end_right;
        if (nums[end_mid] > target) {
            end_right = end_mid - 1;
        } else {
            end_left = end_mid;
        }
    }
        
    ret.push_back(start_left);
    ret.push_back(end_left);
    return ret;
}