[BOJ 7453] 합이 0인 네 정수

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

레벨 : G2
알고리즘 분류 : 정렬, 이분탐색, 투 포인터, 중간에서 만나기

(0) 문제 요약

최대 길이 4000의 네 정수 배열 A, B, C, D를 입력 받고, A[a] + B[b] + C[c] + D[d] = 0이 되는 튜플 (a, b, c, d)의 개수를 구하는 문제.

(1) 풀이 핵심

그냥 for문 4개를 돌려서 가능한 튜플을 구하려고 하면 최대 $4000^4$번의 연산을 해야하므로 TLE에 빠지게 된다.
배열 A, B의 합으로 가능한 것과 C, D의 합으로 가능한 것들을 따로 계산한다면, (일단 여기에서는) $4000^2\times 2$번의 연산만 하면되므로 소요되는 시간을 효과적으로 줄일 수 있다.

(2.1) 풀이 1: map을 이용한 풀이

배열 A, B의 가능한 합들을 key로, 그리고 해당 합이 나온 횟수를 value로 가지는 map(mapAB)을 만들고, 배열 C, D의 가능한 합들도 마찬가지로 map(mapCD)으로 만든다.

이후 mapAB의 임의의 키 k에 대해, mapCD[-k]를 찾아서 정답에 mapAB[k] * mapCD[-k]를 더해주면 된다. (단 mapCD에 키 k가 없다면 0을 곱하도록 해서)

(2.2) 풀이 2: 정렬과 이분 탐색을 이용한 풀이

하지만 풀이 1은 c++에서는 잘 안 된다. map의 경우는 내부 정렬이 일어나서인지 너무 느리고, unordered_map의 경우 메모리 오버헤드가 너무 크기 때문이다.

A, B의 가능한 합들과 C, D의 가능한 합들을 구한 결과 다음과 같은 배열들을 얻었다고 해보자. 위가 A, B의 가능한 합을 담은 배열 sumAB이고, 아래가 C, D의 가능한 합들을 구한 배열 sumCD이다.

앞서 [풀이 1] 에서와 같이 sumAB의 임의의 원소 sumAB[i]에 대해, -sumAB[i]가 sumCD에도 있는지를 확인하면 된다. 물론 이 두 배열의 가능한 합을 구하기 위해서 또 다시 2중 for문을 돌린다면 두 배열을 나눠서 가능한 합을 구한 의미가 없다. 이분 탐색을 이용하도록 하자($O(log N)$)

이분 탐색을 이용하기 위해 sumAB와 sumCD를 모두 오름차순으로 정렬하면 다음과 같다.

그럼 sumAB의 0번째 요소부터 확인해보자. sumAB[0]이 -40이므로 sumCD에 40이 있는지 확인하면 된다.

sumCD에서 40을 찾으려면 sumCD에서 40보다 크거나 같은 가장 작은 수(lower bound)를 찾아보아야 한다

sumCD에서 40보다 크거나 같은 가장 작은 수는 78이다. 즉sumCD에는 40이 없다.

그렇다면 다음 값에 대해 마찬가지의 방법을 적용해보자.

sumAB[1]이 -21이므로 sumCD에서는 21이 있는지 확인하면 된다. sumCD에서 21의 lower bound로서 21을 찾을 수 있다.

문제는 sumAB에는 -21이 2개 있고, sumCD에도 21이 2개가 있다는 것, 즉 가능한 튜플의 수는 4개라는 것이다.

따라서 sumAB에서도 중복된 수가 있다면 그 수의 개수를 모두 세고, sumCD에서도 마찬가지로 해줘야 한다.

sumAB에서 -21의 개수를 세기 위해서는 -21보다 큰 가장 작은 수(upper bound)를 찾으면 된다. 0의 인덱스는 3이고 지금 보고 있는 -21의 인덱스는 1이므로 sumAB에는 총 $3-1 = 2$개의 -21이 있음을 알 수 있다.

sumCD에서도 마찬가지로 21의 개수를 세기 위해 21의 upper bound를 찾아 인덱스를 빼주면 된다. 21의 upper bound는 78, 그 인덱스는 7이고, 21의 lower bound의 인덱스가 5이므로 sumCD에는 총 $7-5=2$개의 21이 있음을 알 수 있다.

정답에 $(3-1)\times (7-5) = 4$를 더해주자. (현재 4)
이후 sumAB에서 나타나는 두 번째 -21은 볼 필요가 없다. 바로 0으로 넘어가서 마찬가지로 계산하자.

sumCD에서 0의 lower bound는 1이니 0이 없다. sumAB에서 0의 upper bound인 2를 확인하자.

sumCD에서 -2의 lower bound는 1이므로 또 넘어간다.

sumCD에서 -4의 lower bound로 -4를 찾아냈다.
sumAB에는 4가 $(4의\,upper\,bound의 \,idx) - (4의\,idx) = (7-6) = 1$개 있고,
sumCD에는 -4가 $(-4의\,upper\,bound의\,idx)-(-4의\,lower\,bound의\,idx) = (2-1) = 1$개가 있음을 알 수 있다. 정답에 $1 \times 1$을 더해주자.(현재 5)

42, 84도 마찬가지로,sumAB에서 더 이상 확인할 게 없을 때까지 위 과정을 진행하면 되고, 위 예시에서 정답은 5가 된다.

3. 소스코드(C++)

아래 코드에서는 위 설명과 달리 sumAB, sumCD가 아니라 sumA, sumB로 쓰여 있다

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

int n, A[4][4000];
long long aIdx, bIdx, nextAIdx, nextBIdx, ans;
vector<int> sumA, sumB;


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

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

    sort(sumA.begin(), sumA.end());
    sort(sumB.begin(), sumB.end());

    while (true){
        if (aIdx == sumA.size()) break;
        bIdx = lower_bound(sumB.begin(), sumB.end(), -sumA[aIdx]) - sumB.begin();
        nextAIdx = upper_bound(sumA.begin(), sumA.end(), sumA[aIdx]) - sumA.begin();
        if (sumA[aIdx] + sumB[bIdx] == 0){
            nextBIdx = upper_bound(sumB.begin(), sumB.end(), sumB[bIdx]) - sumB.begin();
            ans += (nextAIdx - aIdx) * (nextBIdx - bIdx);
        }
        aIdx = nextAIdx;
    }    

    printf("%lld", ans);
}