[BaekJoon 3151번] 합이 0(C++)

※주의※
저의 풀이가 정답은 아닙니다.
다른 코드가 더 효율적이거나 좋을 수 있습니다.
언제나 다른 사람의 코드는 참고만 하시기 바랍니다.

 

[문제 풀이 사이트]

 

3151번: 합이 0

 

[문제 설명]

 

Elly는 예상치 못하게 프로그래밍 대회를 준비하는 학생들을 가르칠 위기에 처했다. 대회는 정확히 3명으로 구성된 팀만 참가가 가능하다. 그러나 그녀가 가르칠 학생들에게는 큰 문제가 있었다. 코딩 실력이 좋으면 팀워크가 떨어지고, 팀워크가 좋을수록 코딩 실력이 떨어진다. 그리고 출전하고자 하는 대회는 코딩 실력과 팀워크 모두가 중요하다.
Elly는 그녀가 가르칠 수 있는 모든 학생들의 코딩 실력을 알고 있다. 각각의 코딩 실력 Ai는 -10000부터 10000 사이의 정수로 표시되어 있다. 그녀는 팀워크와 코딩 실력이 모두 적절한 팀을 만들기 위해, 세 팀원의 코딩 실력의 합이 0이 되는 팀을 만들고자 한다. 이러한 조건 하에, 그녀가 대회에 출전할 수 있는 팀을 얼마나 많이 만들 수 있는지를 계산하여라.
N명의 학생들의 코딩 실력 Ai가 -10000부터 10000사이의 정수로 주어질 때, 합이 0이 되는 3인조를 만들 수 있는 경우의 수를 구하여라.

 

[입력]

 

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

[출력]

 

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

 

[제한]

 

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

 

[입출력 예]

 

입력	출력
10 2 -5 2 3 -4 7 -4 0 1 -6	6

 

[문제 풀이]

 

3명을 선택하고 선택된 3명의 코딩 실력을 합했을 때 0이 되어야 한다.

 

알고리즘을 간단하게 설명하자면

 

1. 2중 for문을 이용하여 2명을 선택한다. 
2. 선택된 두명을 제외한 나머지 인원들에 대해 이분 탐색을 진행한다.

찾고자 하는 값은 선택된 두명의 코딩 실력을 합한 결과에 -1을 곱한 값이다.

 

이분 탐색은 lower_bound와 upper_bound 를 사용하는데

lower_bound의 경우 찾고자 하는 값이 있다면 찾고자 하는 값이 맨 처음으로 나오는 위치를 반환하고

upper_bound의 경우 찾고자 하는 값보다 초과되는 위치를 반환한다.

 

그래서 lower_bound 사용했을 때 범위를 넘어서지 않는다면,

upper_bound 한 위치에서 lower_bound한 위치를 뺀 결과가 찾고자하는 값의 개수가 된다.

 

 

[최종 코드]

 

[GitHub]

 

GitHub - SmallPeanutPark/BAEKJOON: Study BaekJoon for Coding Test

 

#include <bits/stdc++.h>
using namespace std;
vector<int> v;

void input(int N) {
    for(int i = 0; i < N; ++i) {
        int n; cin >> n;
        v.emplace_back(n);
    }
    sort(v.begin(), v.end());
}

long long mix(int N) {
    long long ans = 0;
    for(int first = 0; first < N; ++first) {
        for(int second = first + 1; second < N; ++second) {
            int end = second + 1;
            int lower_location = lower_bound(v.begin() + end, v.end(), -(v[first] + v[second])) - v.begin();
            int upper_location = upper_bound(v.begin() + end, v.end(), -(v[first] + v[second])) - v.begin();
            if(lower_location < N) {
                ans += (upper_location - lower_location);
            }
        }
    }
    return ans;
}

int main(void) {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);
    int N; cin >> N;
    input(N);
    // 세명이 팀이고 세명을 적절히 섞어서 0이 되어야함
    long long answer = mix(N);
    cout << answer;
    return 0;
}