3151. 합이 0

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

3151번: 합이 0

 

풀이1

필요 알고리즘: 이분탐색

 

"숫자 3개를 합한 값이 0이 된다"라는 것은 \(a+b+c=0\)에서 \(a+b=-c\)임을 의미한다. 따라서 배열을 이중 for문으로 순회하며 \(a+b\)를 찾고, \(a+b=-c\)를 만족하는 \(c\)의 값을 나머지 수에서 이분탐색으로 찾으면 된다.

 

탐색에 주의할 점으로, \(a+b=-c\)에서 \(c\)가 여러 개 있을 수 있다. 배열이 {-2, -1, 3, 3}과 같을 경우, \(a=-2\), \(b=-1\)일 때 \(c=3\)인 인덱스가 두 개 존재한다. 만약 binary_search 함수로 값이 존재하는지만을 확인하면 이것이 무시되고 1개만 카운트되므로, upper_bound로 구한 인덱스와 lower_bound로 구한 인덱스의 차로 \(c\)의 개수를 찾아야 한다.

 

또한, 정답은 최대 \({}_{10000}\mathrm{C}_{3}\)가 될 수 있는데, 이는 int 자료형 범위를 넘어서므로 long long 자료형으로 출력해야 하는 것도 잊지 말자.

 

시간 복잡도

C++ intro sort의 시간 복잡도는 \(O(N\log{N})\), \(a+b\)를 찾기 위한 이중 for문의 시간 복잡도는 \(O(N^{2})\), upper_bound와 lower_bound 함수는 모두 이분탐색 알고리즘을 사용하므로 시간 복잡도는 \(O(\log{N^{2}})=O(2\log{N})\)이다. 이를 종합하면 총 시간 복잡도는 \(O(N^{2}\log{N})\)이다.

 

소스코드

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

int N;
long long ans;
int arr[10000];

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    
    cin >> N;
    for(int i = 0; i < N; i++) cin >> arr[i];
    sort(arr, arr + N);
    
    for(int i = 0; i < N - 2; i++) {
        for(int j = i + 1; j < N - 1; j++) {
            int sum = -(arr[i] + arr[j]);
            ans += (upper_bound(arr + j + 1, arr + N, sum) - lower_bound(arr + j + 1, arr + N, sum));
        }
    }
    
    cout << ans;
    return 0;
}

 

풀이2

필요 알고리즘: 투 포인터

 

위의 이진탐색 알고리즘으로 푼 결과 1940ms라는 경이로운 실행 시간이 나왔다. 이걸 더 줄일 수 없을까 찾아보다 투 포인터로도 풀 수 있다는 것을 알 수 있었다(후술할 풀이 과정을 보면 쓰리 포인터라 불러야 할지도 모르겠다).

 

이 풀이의 골자는, \(a+b+c=0\)에서 \(a\)를 하나 찍고, 나머지 범위에서 투 포인터로 \(b\)와 \(c\)를 찍으면서 합이 0이 되는 세 숫자를 찾는 것이다. 탐색할 배열은 사전에 정렬이 필요하다.

 

구체적인 과정을 설명하면 다음과 같다.

1. 배열의 첫 인덱스의 숫자를 \(a\)로 찍는다.
2. \(a\)로 찍은 인덱스의 다음 인덱스의 숫자를 \(b\), 맨 마지막 인덱스의 숫자를 \(c\)로 찍는다.
3. \(a+b+c\)를 계산한다.
   1. 0보다 클 경우, \(c\)를 찍은 인덱스를 1 내린다.
   2. 0보다 작을 경우, \(b\)를 찍은 인덱스를 1 올린다.
   3. 0일 경우, \(b\)와 \(c\)가 같은지 여부에 따라 출력값을 증가시키고, \(b\)를 찍은 인덱스를 올리거나 5로 넘어간다(후술).
4. \(b\)의 인덱스와 \(c\)의 인덱스가 같아질 때까지 3을 반복한다.
5. \(a\)의 인덱스를 다음으로 옮기고, 2~4를 반복한다.
6. 배열의 끝에 도달할 때까지 5를 반복한다.

3에서 \(a+b+c=0\)일 때가 좀 복잡한데, \(b\)와 \(c\)의 값이 어떤지에 따라 과정이 갈린다.

 

-7	-4	2	2	2	3	3	4	4	4
	\(a\)	\(b\)		\(c\)

배열이 있고, 인덱스가 다음과 같이 가리켜 합이 0이 된다고 생각해 보자. \(b=c\)일 경우, 정렬된 배열 특성상 그 사이의 수들도 전부 \(b\), \(c\)와 같다. 따라서 \(a\)는 고정된 상태에서 3개의 2 중 \(b\), \(c\)로 삼을 2개를 뽑는 조합과 같다. 식으로 나타내면 \({}_{n}\mathrm{C}_{2}\)로, 여기서 \(n=(c\mathrm{의\,인덱스})-(b\mathrm{의\,인덱스})+1\)이다.

 

인덱스 사이의 값이 전부 같은 이상, 탐색을 계속하는 의미가 없으므로 여기서 중지하고 다음 \(a\)로 넘어가 다시 탐색을 실시한다.

 

-7	-4	2	2	2	3	3	4	4	4
\(a\)					\(b\)				\(c\)

이번엔 인덱스가 다음과 같이 가리켜서 합이 0이 됐다고 생각해 보자. \(b\neq c\)일 경우, \(b\) 인덱스 이후와 \(c\) 인덱스 이전 각각에 같은 수가 존재하는지 확인해야 한다. \(b\)가 2개, \(c\)가 3개 존재하므로 경우의 수는 (\(b\)의 개수) × (\(c\)의 개수)가 된다.

 

이 경우는 인덱스 사이에 만족하는 값이 또 존재할 수 있으므로, (\(b\)의 개수)만큼 인덱스를 이동시키고 탐색을 계속한다.

 

시간 복잡도

\(N\)개의 \(a\)에 대하여 투 포인터 알고리즘을 적용하므로, 시간 복잡도는 \(O(N^{2})\)이다.

 

소스코드

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

int N;
long long ans;
int arr[10000];

long long getComb(int n) {
    return n * (n - 1) / 2;
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    
    cin >> N;
    for(int i = 0; i < N; i++) cin >> arr[i];
    sort(arr, arr + N);
    
    for(int i = 0; i < N - 2; i++) {
        int first = i + 1, last = N - 1;
        while(first < last) {
            int sum = arr[i] + arr[first] + arr[last];
            if(sum == 0) {
                if(arr[first] == arr[last]) {
                    ans += getComb(last - first + 1);
                    break;
                }
                else {
                    int left = 0, right = 0;
                    while(arr[first + left] == arr[first + left + 1]) left++;
                    while(arr[last - right] == arr[last - right - 1]) right++;
                    ans += ((left + 1) * (right + 1));
                    first += (left + 1);
                }
            }
            else if(sum > 0) last--;
            else first++;
        }
    }
    
    cout << ans;
    return 0;
}