[C++] 프로그래머스 2의 영역

문제설명

정수 배열 arr가 주어집니다.
배열 안의 2가 모두 포함된 가장 작은 연속된 부분 배열을
return 하는 solution 함수를 완성해 주세요.
단, arr에 2가 없는 경우 [-1]을 return 합니다.

제한사항

• 1 ≤ arr의 길이 ≤ 100,000
  ◦ 1 ≤ arr의 원소 ≤ 10

입출력 예

arr	result
[1, 2, 1, 4, 5, 2, 9]	[2, 1, 4, 5, 2]
[1, 2, 1]	[2]
[1, 1, 1]	[-1]
[1, 2, 1, 2, 1, 10, 2, 1]	[2, 1, 2, 1, 10, 2]

풀이

arr 배열에 포함된 2를 찾아 loc라는 배열을 새로 만들었고, loc 배열을 통해 문제 조건에 맞는지 판단하여 리턴값을 리턴하는 로직을 구현했습니다.

#include <string>
#include <vector>

using namespace std;

vector<int> solution(vector<int> arr) {
    vector<int> answer;
    vector<int> loc;
    
    for(int i = 0; i < arr.size(); i++)
    {
        if(arr[i] == 2)
            loc.push_back(i);
    }
    
    if(loc.size() == 0)
        answer.push_back(-1);
    else if(loc.size() == 1)
        answer.push_back(arr[loc[0]]);
    else if(loc.size() >= 2)
    {
        for(int i = loc[0]; i <= loc.back(); i++)
        {
            answer.push_back(arr[i]);
        }
    }
    
    return answer;
}

결과