[BOJ 7435] 합이 0인 네 정수

문제는 여기서 풀 수 있습니다.

문제 요약

정수로 이루어진 배열 네 개가 있습니다. 이 때 각각의 배열에서 한 개의 수를 뽑아 더한 결과가 0이 되는 경우의 수를 찾는 문제입니다.

문제 풀이

가장 쉽게 푸는 방법은 완전 탐색입니다. 완전 탐색은 네 개의 배열에서 수를 하나씩 선택할 수 있는 모든 경우의 수를 확인해야 해서 $O(N^4)$의 시간 복잡도를 가집니다. 하지만 문제에서 주어지는 N의 최대 크기 때문에 $O(N^4)$의 시간 복잡도는 시간제한을 통과하지 못합니다.

만약 네 개의 배열을 두 개의 배열로 줄일 수 있다면 어떨까요? 주어진 네 개의 배열 중 두 개인 $A$와 $B$로 만들 수 있는 합을 전부 저장한 배열 $X$와 나머지 두 배열인 $C$와 $D$로 만들 수 있는 합을 전부 저장한 배열 $Y$가 있습니다. 이때 $X$에서 수 하나를 선택하고 $Y$에서 수 하나를 선택한 뒤 더한 값이 0이라면 문제에서 주어진 네 개의 배열에서 수를 하나씩 선택해 더한 결과가 0인 것과 같습니다.

이제 남은 일은 두 개의 배열에서 합이 0이 되는 모든 경우를 찾아봅시다. 만약 $X$에서 선택한 수가 10이라면 $Y$에서 선택한 수가 -10 이어야 두 수의 합이 0이 됩니다. 즉 $X$에서 선택한 수의 반대 부호를 가진 수가 $Y$에 몇 개 있는지 안다면 문제를 풀 수 있습니다.

이 작업을 빠르게 하기 위해 정렬된 배열에서 원하는 값을 찾는 $O(logN)$에 찾을 수 있는 알고리즘인 이진 검색을 이용하면 $O(N^2logN^2)$에 푼제를 풀 수 있습니다.

#include <stdio.h>
#include <vector>
#include <algorithm>

using namespace std;

int n, A[4010][4];
vector<int> X, Y;

int main() {
    scanf("%d", &n);
    for (int i = 0; i < n; ++i) {
        for (int j = 0; j < 4; ++j) {
            scanf("%d", &A[i][j]);
        }
    }

    for (int i = 0; i < n; ++i) {
        for (int j = 0; j < n; ++j) {
            X.push_back(A[i][0] + A[j][1]);
            Y.push_back(A[i][2] + A[j][3]);
        }
    }

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

    long long ans = 0;
    for (int val : X) {
        int target = -val;
        auto upIt = upper_bound(Y.begin(), Y.end(), target);
        auto loIt = lower_bound(Y.begin(), Y.end(), target);
        ans += upIt - loIt;
    }

    printf("%lld\n", ans);
}