문제 설명
정수로 이루어진 크기가 같은 배열 A, B, C, D가 있다.
A[a], B[b], C[c], D[d]의 합이 0인 (a, b, c, d) 쌍의 개수를 구하는 프로그램을 작성하시오.
https://www.acmicpc.net/problem/7453
제한 사항
풀이
문제를 요약하면, A, B, C, D 4개의 배열이 주어지면 각 하나의 요소를 선택해 그 합이 0을 만드는 경우의 수를 찾아야 한다.
문제는 간단하지만, 4개의 배열을 살펴봐야 하기 때문에 $O(N^4)$이 된다.
따라서, 이를 해결하기 위해 약간의 변형이 필요하다.
처음에는 비슷한 유형을 풀어본 적이 있어서 A, B를 2중 for문을 통해 정하고 나머지를 투 포인터로 해결할 수 있다고 생각했다.
하지만, 남은 C, D의 포인터를 움직이는 규칙을 만들어 낼 수 없었다.
이 문제를 풀기 위해서는 두 배열씩 짝을 지어 임시 결과 배열을 생성한 뒤 정렬해야 한다.
어떻게 짝을 지어도 상관없지만, AB, CD를 묶는 게 보기 좋으니 이렇게 진행했다.
그리고 AB를 기준으로 CD에서 -AB[i]를 찾으면 합으로 0을 만들 수 있다.
CD에는 -AB[i]가 여러 개 존재할 수 있다.
따라서, 가장 앞에 있는 인덱스와 -AB[i]보다 큰 가장 빠른 인덱스의 차이를 계산하면 그 개수를 구할 수 있다.
여기에 가장 적합한 함수는 lower_bound, upper_bound이다.
- lower_bound(n): n이상 중 가장 빠른 요소
- upper_bound(n): n초과 중 가장 빠른 요소
vector<long long> AB;
vector<long long> CD;
for (int i = 0; i < N; i++)
{
for (int j = 0; j < N; j++)
{
AB.push_back(A[i] + B[j]);
CD.push_back(C[i] + D[j]);
}
}
sort(AB.begin(), AB.end());
sort(CD.begin(), CD.end());
int ans = 0;
for (int i = 0; i < AB.size(); i++)
{
auto under = lower_bound(CD.begin(), CD.end(), -AB[i]) - CD.begin();
auto over = upper_bound(CD.begin(), CD.end(), -AB[i]) - CD.begin();
ans += (over - under);
}
전체 코드
#include <bits/stdc++.h>
using namespace std;
int N;
int main() {
ios_base::sync_with_stdio(0);
cin.tie(0);
cout.tie(0);
cout << fixed;
cout.precision(4);
cin >> N;
vector<long long> A;
vector<long long> B;
vector<long long> C;
vector<long long> D;
for (int i = 0; i < N; i++)
{
long long a, b, c, d;
cin >> a >> b >> c >> d;
A.push_back(a);
B.push_back(b);
C.push_back(c);
D.push_back(d);
}
vector<long long> AB;
vector<long long> CD;
for (int i = 0; i < N; i++)
{
for (int j = 0; j < N; j++)
{
AB.push_back(A[i] + B[j]);
CD.push_back(C[i] + D[j]);
}
}
sort(AB.begin(), AB.end());
sort(CD.begin(), CD.end());
int ans = 0;
for (int i = 0; i < AB.size(); i++)
{
auto under = lower_bound(CD.begin(), CD.end(), -AB[i]) - CD.begin();
auto over = upper_bound(CD.begin(), CD.end(), -AB[i]) - CD.begin();
ans += (over - under);
}
cout << ans;
return 0;
}