백준 3151 - 합이 0

문제 설명

Elly는 예상치 못하게 프로그래밍 대회를 준비하는 학생들을 가르칠 위기에 처했다. 대회는 정확히 3명으로 구성된 팀만 참가가 가능하다. 그러나 그녀가 가르칠 학생들에게는 큰 문제가 있었다. 코딩 실력이 좋으면 팀워크가 떨어지고, 팀워크가 좋을수록 코딩 실력이 떨어진다. 그리고 출전하고자 하는 대회는 코딩 실력과 팀워크 모두가 중요하다.

Elly는 그녀가 가르칠 수 있는 모든 학생들의 코딩 실력을 알고 있다. 각각의 코딩 실력 Ai는 -10000부터 10000 사이의 정수로 표시되어 있다. 그녀는 팀워크와 코딩 실력이 모두 적절한 팀을 만들기 위해, 세 팀원의 코딩 실력의 합이 0이 되는 팀을 만들고자 한다. 이러한 조건 하에, 그녀가 대회에 출전할 수 있는 팀을 얼마나 많이 만들 수 있는지를 계산하여라.

N명의 학생들의 코딩 실력 Ai가 -10000부터 10000사이의 정수로 주어질 때, 합이 0이 되는 3인조를 만들 수 있는 경우의 수를 구하여라.

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

---

 

 

 

 

 

 

제한 사항

---

 

 

 

 

 

 

풀이

문제를 요약하면, 주어진 수들 중 세 개를 뽑아 합이 0이 되도록하는 조합의 수를 구하면 된다.

 

간단하게 생각해보면 각각 하나씩 살펴보며 세 개의 수를 뽑아 0이 되는지 확인하면 된다.

하지만, 이렇게 한다면 $O(N^3)$으로 시간 초과가 발생한다.

 

이를 줄이려면 한 가지 아이디어가 필요하다.

0으로 만들기 위해서 두 개의 수를 고른다면 나머지 수가 결정된다.

그럼 $O(N^2)$로 시간을 줄일 수 있다.

하지만, 해당 문제에서는 중복된 수가 존재하기 때문에 나머지 수가 여러 개일 수 있다.

이를 처리하기 위해서는 이진 탐색을 적절히 활용하면 된다.

우선, 모든 수를 정렬한 뒤 나머지 수를 lower_bound를 통해 찾아낸다.

그리고 같은 수의 마지막을 찾기 위해 upper_bound를 통해 마지막을 찾아낸다.

이 두 위치를 뺀 값이 유효한 조합의 개수이다.

a, b가 정해졌다면 찾아야 하는 수는 0 - (a + b)이고 이를 이진탐색으로 시작과 끝을 찾아내는 것이다.

여기서 주의해야 할 점은 모든 범위를 탐색하면 a, b를 포함할 수 있다.

따라서, b(오른쪽)의 수보다 큰 수에 대해서만 진행해야 한다.

---

 

 

 

 

 

전체 코드

#include <bits/stdc++.h>
using namespace std;
#define INPUT_OPTIMIZE cin.tie(NULL); cout.tie(NULL); ios::sync_with_stdio(false);
#define INF 2e9

using namespace std;
int N;

vector<int> abilities;

int main() 
{
	INPUT_OPTIMIZE;

	cin >> N;
	abilities.resize(N);

	for (int i = 0; i < N; i++)
	{
		cin >> abilities[i];
	}

	sort(abilities.begin(), abilities.end());

	/*
	* 3개의 합이 0이 되는 조합을 많이 만들기
	*/

	long long ans = 0;

	for (int i = 0; i < N; i++)
	{
		for (int j = i+1; j < N; j++)
		{
			int a = abilities[i];
			int b = abilities[j];

			int sum = a + b;

			int pos = lower_bound(abilities.begin() + j + 1, abilities.end(), -sum) - abilities.begin();
			int pos2 = upper_bound(abilities.begin() + j + 1, abilities.end(), -sum) - abilities.begin();

			if (pos < abilities.size() && abilities[pos] == -sum)
			{
				ans += pos2 - pos;
			}
		}
	}

	cout << ans;

	return 0;
}

---