[백준 7453] 합이 0인 네 정수 (C++)

https://www.acmicpc.net/problem/7453

문제

---

정수로 이루어진 크기가 같은 배열 A, B, C, D가 있다.

A[a], B[b], C[c], D[d]의 합이 0인 (a, b, c, d) 쌍의 개수를 구하는 프로그램을 작성하시오.

 

입력

---

첫째 줄에 배열의 크기 n (1 ≤ n ≤ 4000)이 주어진다. 다음 n개 줄에는 A, B, C, D에 포함되는 정수가 공백으로 구분되어져서 주어진다. 배열에 들어있는 정수의 절댓값은 최대 2²⁸이다.

 

출력

---

합이 0이 되는 쌍의 개수를 출력한다.

 

문제 풀이

---

MITM 알고리즘과 투포인터를 같이 사용하는 문제.

A, B, C, D 4개 배열을 모든 A+B를 저장하는 벡터 v1, 모든 C+D를 저장하는 벡터 v2 2개로 만들어 시간 복잡도를 줄인다. 벡터를 사용해 2중 for문으로 모든 A[i]+B[j], C[i]+D[j]를 구한 후, 각각 정렬하여 투포인터로 쌍의 개수를 구한다. v1과 v2를 <두 수의 합, 나온 횟수>를 저장하는 map 형태로 구현할 시, 시간초과가 발생하므로 벡터로 구현하였다.

 

이제 v1를 탐색하는 인덱스를 start, v2를 탐색하는 인덱스를 end로 잡고, v1[start]와 v2[end]를 더했을 때 0이 되는 부분과, 그 쌍의 개수를 찾을 것이다. start는 v1의 시작지점부터 시작하여 점차 증가하고, end는 v2의 끝 부분부터 시작하여 점차 감소하는 형태로 구현한다.

start와 end가 벡터의 범위를 벗어나기 전까지, 현재 v1[start]와 v2[end]의 합이 0보다 작거나 같을 때까지 end를 감소시키고, 만약 v1[start]와 v2[end]의 합이 0이라면 같은 값을 가지는 쌍이몇 쌍이 존재할 수 있는지 확인한다. 이후 start를 1 증가시킨다.

 

v1[start]와 v2[end]의 합이 0일 때, 같은 값을 가지는 쌍을 구하는 방법은 v1에서 start를 증가시켜가며 이전 값과 동일한 값이 몇 개 있는지 살핀 후, v2에서 end를 감소시켜가며 이전과 동일한 값이 몇개 있는지 확인 후 이 둘을 곱해주면 된다.

while문의 끝에서 start가 1 증가하므로, 이전 값과 동일 한 값이 아니라면 start를 1 감소시켜 while문의 끝에서 제대로 다음 index를 가리킬 수 있게 한다.

 

아래는 코드.

#include <iostream>
#include <algorithm>
#include <vector>
using namespace std;

int main()
{
	cin.tie(NULL);
	ios::sync_with_stdio(false);

	int N;
	cin >> N;
	int* A = new int[N];
	int* B = new int[N];
	int* C = new int[N];
	int* D = new int[N];
	vector<int> v1;
	vector<int> v2;
	long long answer = 0;
	for (int i = 0; i < N; i++)
	{
		cin >> A[i] >> B[i] >> C[i] >> D[i];
	}
	for (int i = 0; i < N; i++)
	{
		for (int j = 0; j < N; j++)
		{
			v1.push_back(A[i] + B[j]);
			v2.push_back(C[i] + D[j]);
		}
	}
	sort(v1.begin(), v1.end());
	sort(v2.begin(), v2.end());
	int start = 0;
	int end = v2.size() - 1;
	while (start < v1.size() && end>=0)
	{
		while (end >= 0 && v1[start] + v2[end] > 0)
		{
			end--;
		}
		if (end >= 0 && v1[start] + v2[end] == 0)
		{
			long long aCount = 1;
			long long bCount = 1;
			while (start < v1.size())
			{
				start++;
				if (start< v1.size() && v1[start] == v1[start - 1])
				{
					aCount++;
				}
				else
				{
					start--;
					break;
				}
			}
			while (end >= 0)
			{
				end--;
				if (end>=0 && v2[end] == v2[end + 1])
				{
					bCount++;
				}
				else
				{
					break;
				}
			}
			answer += aCount * bCount;
		}
		start++;
	}
	cout << answer << "\n";
	return 0;
}