백준 3151번: 합이 0

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

3151번: 합이 0

이분 탐색과 투 포인터, 동적 계획법 세 가지 방법으로 해결할 수 있다.

• 이분 탐색

3인조의 코딩 실력의 합이 0일 경우, i번째 학생의 코딩 실력이 A_i일 때 나머지 두 학생의 코딩 실력의 합은 -A_i이다. 따라서 코딩 실력의 합이 k인 임의의 2인조를 구한 뒤, 코딩 실력이 -k인 학생의 수를 이분 탐색을 이용해 계산한다면 제한 시간 안에 코딩 실력의 합이 0이 되는 3인조의 수를 구할 수 있다.

#include <algorithm>
#include <cstdio>

using namespace std;

#define MAX_N 10000
#define ARG A + j + 1, A + N, -(A[i] + A[j])

int main() {
    // N: 학생의 수, A[i]: i번째 학생의 코딩 실력
    int N, A[MAX_N];
    // 합이 0이 되는 3인조의 수
    unsigned long long ans = 0;

    // 문제의 조건을 입력받은 뒤
    scanf("%d", &N);
    for (int i = 0; i < N; i++) scanf("%d", A + i);
    // 이분 탐색을 위해 코딩 실력을 정렬
    sort(A, A + N);

    // 만들 수 있는 모든 2인조에 대해 
    for (int i = 0; i < N - 2 && A[i] <= 0; i++)
        for (int j = i + 1; j < N - 1; j++)
            // 이분 탐색을 통해 합을 0으로 만드는 3번째 학생의 수를 ans에 더한다
            ans += upper_bound(ARG) - lower_bound(ARG);

    // 코딩 실력의 합이 0이 되는 3인조의 수를 출력
    printf("%lld\n", ans);

    return 0;
}

• 투 포인터

투 포인터를 이용할 경우, 코딩 실력을 정렬한 뒤 학생 하나를 골라 투 포인터를 실행하면 된다. 이 때 코딩 실력이 같은 학생이 여러 명 있을 수 있으므로, 코딩 실력의 합이 0인 3인조를 발견한 경우 코딩 실력이 같은 학생들을 빠르게 넘겨야 한다.

#include <algorithm>
#include <cstdio>

using namespace std;

typedef unsigned long long ull;

#define MAX_N 10000
#define SUM (A[i] + A[l] + A[r])

int main() {
    // N: 학생의 수, A[i]: i번째 학생의 코딩 실력
    // l, r: 투 포인터에 사용할 인덱스
    // dl, dr: l과 r의 변위
    int N, A[MAX_N], l, r, dl, dr;
    // 합이 0이 되는 3인조의 수
    ull ans = 0;

    // 문제의 조건을 입력받은 뒤
    scanf("%d", &N);
    for (int i = 0; i < N; i++) scanf("%d", A + i);
    // 투 포인터를 위해 코딩 실력을 정렬
    sort(A, A + N);

    // 3인조에서 코딩 실력이 가장 작은 학생을 먼저 고른 뒤
    for (int i = 0; i < N - 2 && A[i] <= 0; i++) {
        // 투 포인터의 시작점을 l과 r에 저장
        l = i + 1, r = N - 1;
        // 투 포인터 실행
        while (l < r) {
            // 만일 3인조의 코딩 실력의 합이 0이라면
            if (!SUM) {
                // l번째 학생과 r번째 학생의 코딩 실력이 같다면
                if (A[l] == A[r]) {
                    // 코딩 실력이 같은 학생의 수는 (r - l + 1)명이다
                    // 해당 학생들을 순서에 상관없이 두 명을 고른 경우의 수를 ans에 더한 뒤
                    ans += (r - l) * (r - l + 1) / 2;
                    // 이 이상 진행할 필요가 없으므로 break
                    break;
                }
                // l번째 학생과 r번째 학생의 코딩 실력이 다르다면
                else {
                    // l과 r의 변위를 0으로 초기화한 뒤
                    dl = dr = 0;
                    // l번째 학생, r번째 학생과 코딩 실력이 같은 학생 수를 각각 dl, dr에 저장
                    do dl++; while (A[l + dl] == A[l]);
                    do dr++; while (A[r - dr] == A[r]);
                    // ans에 두 학생 수의 곱을 더한 뒤
                    ans += dl * dr;
                    // l과 r을 각각 그 변위만큼 점프
                    l += dl;
                    r -= dr;
                }
            }
            // 3인조의 코딩 실력의 합이 0이 아니라면 투 포인터 진행
            else SUM < 0 ? l++ : r--;
        }
    }

    // 코딩 실력의 합이 0이 되는 3인조의 수를 출력
    printf("%lld\n", ans);

    return 0;
}

• 동적 계획법

3인조의 코딩 실력의 합이 0일 경우, i번째 학생의 코딩 실력이 A_i일 때 나머지 두 학생의 코딩 실력의 합은 -A_i이다. 따라서 임의의 두 학생의 코딩 실력의 합이 k가 되는 경우의 수를 모두 계산한 뒤, 다른 학생의 코딩 실력 A_i를 입력받을 때마다 기존에 구한 2인조의 코딩 실력이 -A_i인 경우의 수를 정답에 더하고, 입력받은 학생을 포함하는 2인조를 만들어 저장한다. 이 때 메모이제이션을 한 뒤 ans를 갱신할 경우, i번째 학생을 3인조에 두 번 포함시킬 가능성이 있으므로 ans를 갱신한 뒤 메모이제이션을 해야 한다.

#include <cstdio>

typedef unsigned long long ull;

#define MAX_N 10000
#define MAX_SUM 20000

int main() {
    // N: 학생의 수, A[i]: i번째 학생의 코딩 실력
    int N, A[MAX_N];
    // ans: 합이 0이 되는 3인조의 수, dp[MAX_SUM - k]: 합이 k가 되는 2인조의 수
    ull ans = 0, dp[MAX_SUM * 2 + 1] = { 0, };
    
    // 학생의 수를 입력받은 뒤 각 학생에 대해
    scanf("%d", &N);
    for (int i = 0; i < N; i++) {
        // 현재 학생의 코딩 실력을 입력받고
        scanf("%d", A + i);
        // 코딩 실력의 합이 현재 학생의 실력의 -1배가 되는 2인조의 수를 ans에 더한 뒤
        ans += dp[MAX_SUM - A[i]];
        // 현재 학생을 포함하는 2인조를 dp에 저장
        for (int j = 0; j < i; j++) dp[MAX_SUM + A[i] + A[j]]++;
    }

    // 코딩 실력의 합이 0이 되는 3인조의 수를 출력
    printf("%lld\n", ans);

    return 0;
}