문제
풀이
누적 합과 이분 탐색을 사용한다.
- 피자 조각의 최대 개수는 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;
}
'알고리즘 > 백준' 카테고리의 다른 글
[백준][C++] 11559: Puyo Puyo (1) | 2023.11.09 |
---|---|
[백준][C++] 20056: 마법사 상어와 파이어볼 (3) | 2023.11.09 |
[백준][C++] 1561: 놀이 공원 (0) | 2023.11.06 |
[백준][C++] 2352: 반도체 설계 (0) | 2023.11.05 |
[백준][C++] 3079: 입국심사 (0) | 2023.11.04 |