Lesson 5. Perfix Sums - PassingCars #4 (번외-prefix sums)

이 문제는 (PassingCars #3)가 가장 좋은 코드다.

여기서는 번외로, 부분합(prefix sums) 해법을 제시한다.

 

1. 코드 (C++)

int solution(vector<int> &A) {

    vector<int> P = {}; //0, 동쪽으로 이동하는 차량
    int q_cnt = 0; //서쪽으로 이동하는 차량 수
    
    //동쪽으로 이동하는 차량을 vector P에 저장하고,
    //서쪽으로 이동하는 차량의 수를 구한다.
    for(size_t i = 0; i < A.size(); i++)
    {
        if(A[i] == 0) P.emplace_back(i);
        else q_cnt++;
    }

    //서쪽으로 이동하는 차량의 부분합(prefix sums) 배열을 정의한다.
    vector<int> PREFIX_SUMS_Q(A.size(), 0);
    PREFIX_SUMS_Q[0] = q_cnt;
    if(A[0] == 1) q_cnt--;
    
    //prefix sums 배열에 값을 채운다. 값은 서쪽으로 이동하는 차량의 수이다.
    for(size_t i = 1; i < A.size(); i++)
    {
        PREFIX_SUMS_Q[i] = q_cnt;
        if(A[i] == 1) 
        {
            q_cnt--;
            if(q_cnt <= 0) break;
        }
    }

    //동쪽으로 이동하는 차량이, 서쪽으로 이동하는 차량을 만나는 대수를
    //부분합(prefix sums) 배열을 이용해서 계산한다.
    int cnt = 0;
    for(size_t i = 0; i < P.size(); i++)
    {
        int pos = P[i];
        cnt += PREFIX_SUMS_Q[pos];
        if(cnt > 1000000000) return -1;
    }    

    return cnt;
}

 

 

2. 결과