백준 3151번 합이 0 / C++ / 투포인터 (Tow Pointer)

알고리즘 개념: 투포인터

투포인터에 대해 잘 모른다면 이전에 정리해 둔 개념 정리를 한번 읽어보고 오는 것을 추천합니다. 

알고리즘 공부 2주차: 투포인터 개념

알공 2주차: 투포인터 개념

백준 3151번 합이 0

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

3151번: 합이 0

문제

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

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

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

입력

입력은 표준 입력으로 주어진다.

첫 번째 줄에 학생의 수 N이 입력된다. 두 번째 줄에 N개의 그녀가 가르칠 학생들의 코딩 실력인 Ai가 주어진다.

출력

표준 출력으로 그녀가 고를 수 있는 팀의 수를 하나의 정수로 출력한다.

제한

• 1 ≤ N ≤ 10000
• -10000 ≤ Ai ≤ 10000

예제 보려면 더보기 클릭

예제 입력 1

10 
2 -5 2 3 -4 7 -4 0 1 -6

예제 출력 1

6

---

문제 설명

문제 자체는 정말 간단하다. 3명의 합이 0이 되는 모든 경우의 수를 구하면 되는 것이다.

그러나 어쩌면 당연하게도 이것을 하나하나 계산해가며 반복하면 시간초과가 뜬다.. (이미 해봤다.)

for (ll i = 0; i < n; i++) {
		for (ll j = i+1; j <n; j++) {
			for (ll k = j+1; k < n; k++) {
				if (i != j && j != k && k != i) {
					if (v[i] + v[j] + v[k] == 0) {
						cnt++;
						//cout << v[i] << " " << v[j] << " " << v[k] << "\n";
					
					}
					
				}
			}
		}
	}

각각의 시간복잡도가 n만큼 작용하므로 O(n x n x n)해서 O(n^3)만큼의 시간이 걸린다.. 

n의 최댓값이 10000이니 10000^3을 하면... 왜 시간초과가 나오는지 알 것 같다. 

아무튼 일일이 접근해서 반복하는 것이 안된다면 투포인터를 통해 접근해보자. 

처음 for문은 그대로 두어서 i가 반복되도록 하고 각각 처음과 끝 포인터를 움직여가며 모든 경우의 수를 비교해보는 것이다. 

처음 start부분의 포인터는 s=i+1 (이미 i는 존재하니까)로 두고 끝부분 end는 e=n-1로 시작한다. 

그리고 모든 수의 합이 0이여야 할때 예시를 보면

2 -5 2 3 -4 7 -4 0 1 -6

i=0일때 v[i]는 2이고, v[s]=-5, v[e]=-6이다. 

v[i]+v[s]+v[e]=-9이므로 0이 아닌 것은 당연하지만 이걸 다르게 표현해보면

v[i]+v[e]=-9-v[s] 이다.

 만약 v[i]+v[e]==-v[s]이면 모두 더해서 0인 것이고, v[s]보다 크다면 end를 움직이고 작다면 start를 움직인다. 

if (v[i] + v[e] > -v[s]) e--;
else s++;

처음엔 이걸 이용해서 코드를 짰는데 반례가 있었다. 

바로 코딩 실력이 같은 참가자들이 여러명이 있을 수 있단 것이다. 

예시를 들어보면

8
-2 1 1 1 1 1 1 1

answer : 21

i=0, s=1, e=6일때 -2 , 1 , 1이다.

이때 v[s]와 v[e]과 같은것 만으로 전체가 같다고 가정할 수 있는데 1의 개수가 7개이므로 전체 경우의 수는 7C1이다. 

7C1은 7x(7-1)/2=21이며 1의 개수는 e-s+1을 통해 구할 수 있다. 

코드로 작성하면 다음과 같다. 

if (v[e] == v[s]) {
	ll b = (e - s + 1);
	cnt += b * (b - 1) / 2;
}

그럼 앞선 반례는 해결이 되었다. 

이제 두번째 반례를 보자

10
-1 0 0 0 0 0 1 1 1 1

 

위 예시에서의 경우의 수는 앞선 조합이 아닌 곱셈이다. 

i=0, s=1, e=9인 경우를 보자. v[i]=-1, v[s]=0, v[e]=9 

각각 v[s]와 v[e]는 다르지만 현재 v[i]+v[s]+v[e]=0이다.

v[s]는 0이므로 v[s]에 들어갈 수 있는 0의 개수도 5개이다. 

v[e]에 들어갈 수 있는 1의 개수도 4개이다. 

각각 연속으로 같다면 다를때까지 카운트 개수를 세는 것이다.  

총 경우의 수는 v[i]는 고정이므로 5x4해서 20개이다. 

정리하자면 cnt += v[s]갯수 * v[e]갯수로 나타낼 수 있다. 

while (v[s] == v[s + 1] && s + 1 < e) {
	h++;//v[s]의 갯수
	s++;
}
while (v[e] == v[e - 1] && s < e - 1) {
	k++;//v[e]의 갯수
	e--;
}
	cnt += h * k;

위 모든 걸 정리하면 제출한 코드가 된다. 

---

제출한 코드(C++)

#include <iostream>
#include <string>
#include <algorithm>
#include <vector>
using namespace std;
using ll = long long int;
int main(void) {
	ios::sync_with_stdio(false);
	cin.tie(0);
	cout.tie(0);
	ll n;
	cin >> n;
	int* v = new int[n + 5];
	for (int i = 0; i < n; i++) {
		cin >> v[i];
	}
	ll cnt = 0;
	sort(v, v + n);
	//vector <ll> a;
	ll s, e = n - 1;
	ll h = 1, k = 1;
	ll b;
	for (ll i = 0; i < n; i++) {
		e = n - 1;
		s = i + 1;
		while (s < e) {
			h = 1, k = 1;
			if (v[i] + v[e] == -v[s] && s != i && s < e && i != e) {
				if (v[e] == v[s]) {
					b = (e - s + 1);
					//-10 5 5 5 5 5 5 5 5 과 같은 경우
					//경우의 수 계산
					cnt += b * (b - 1) / 2;
					break;
				}
				else {
					while (v[s] == v[s + 1] && s + 1 < e) {
						h++;
						s++;
					}
					while (v[e] == v[e - 1] && s < e - 1) {
						k++;
						e--;
					}
					cnt += h * k;
				}

			}
			if (v[i] + v[e] > -v[s]) e--;
			else s++;
			if (v[i] > 0 && v[e] > 0 || v[i] < 0 && v[e] < 0) break;
		}



	}
	cout << cnt;

}