[백준][C++] 2632: 피자 판매

문제

풀이

누적 합과 이분 탐색을 사용한다.

• 피자 조각의 최대 개수는 1000개 이므로 O(N * N)으로 해도 시간초과가 발생하지 않는다. 
• A와 B에서 얻을 수 있는 모든 피자 조각을 계산한다. 
• 피자 조각은 끝과 처음이 연결될 수 있으므로 유의해서 더해야 한다. 
• A에서 얻을 수 있는 피자 조각을 기준으로 P와의 차이만큼의 B의 피자 조각의 수를 이분 탐색으로 구한다.

코드

#include <iostream>
#include <vector>
#include <algorithm>

using namespace std;

vector<int> A;
vector<int> A_parts;
vector<int> B;
vector<int> B_parts;

int main()
{
    ios_base::sync_with_stdio(0);
    cin.tie(0);
    cout.tie(nullptr);

    int P;
    int m, n;
    cin >> P >> m >> n;

    int sum_A = 0;
    for (int i = 0; i < m; ++i)
    {
        int temp;
        cin >> temp;
        A.push_back(temp);
        
        sum_A += temp;
    }
    int sum_B = 0;
    for (int i = 0; i < n; ++i)
    {
        int temp;
        cin >> temp;
        B.push_back(temp);
        
        sum_B += temp;
    }

    // 3중 for문은 시간초과
    // A의 i부터 (i + j) 조각까지의 합을 A_parts에 추가한다.
    // 0개와 m개의 합은 중복될 수 있으니 따로 추가한다.
    for (int i = 0; i < m; ++i)
    {
        int temp = 0;
        for (int j = 0; j < m - 1; ++j)
        {
            temp += A[(i + j) % m];
            A_parts.push_back(temp);
        }
    }
    A_parts.push_back(0);
    A_parts.push_back(sum_A);
    
    // B의 i부터 (i + j) 조각까지의 합을 B_parts에 추가한다.
    // 0개와 m개의 합은 중복될 수 있으니 따로 추가한다.
    for (int i = 0; i < n; ++i)
    {
        int temp = 0;
        for (int j = 0; j < n - 1; ++j)
        {
            temp += B[(i + j) % n];
            B_parts.push_back(temp);
        }
    }
    B_parts.push_back(0);
    B_parts.push_back(sum_B);

    // 이진 탐색을 위한 정렬
    sort(B_parts.begin(), B_parts.end());

    // upper_bound, lower_bound로 A_parts[i]와 맞는 B_parts의 수를 찾는다.
    int answer = 0;
    for (int i = 0; i < A_parts.size(); ++i)
    {
        int diff = P - A_parts[i];
        answer += upper_bound(B_parts.begin(), B_parts.end(), diff)
            - lower_bound(B_parts.begin(), B_parts.end(), diff);
    }

    cout << answer; 

    return 0;
}