[백준] 7453번 합이 0인 네 정수 (C++)

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

 

 

이 문제는 배열 A, B, C, D개에서 각각 N개만큼 원소가 주어졌을때 

A[i] + B[j] + C[k] + D[l] = 0 일때 i, j, k, l이 될 수 있는 쌍의 개수를 구하는 문제이다.

 

N이 4000이라 완전탐색을 할 경우 O(N^4)으로 어마어마한 수가 나온다. 이를 최적화 하는 방법을 생각해보자

 

4개의 배열을 합쳐서 배열을 줄이는 방법을 생각해보면 A, B 배열을 a라는 배열과 C, D 배열을 b라는 배열로 새로 만들어보자

 

A, B 배열이 나올 수 있는 모든 쌍을 a라는 배열에 넣고 C, D배열이 나올 수 있는 모든 쌍을 b라는 배열에 넣으면

크기는 2*N^2 으로 1억이 넘지 않은 숫자가 된다.

 

이제 2개의 배열이 되었으니 a배열의 하나 인덱스를 잡고 b배열의 하나 인덱스를 잡아서 0이 나오는 경우를 찾아보면 답이 나올 것이다.

 

이때 배열 하나가 크기 때문에 정렬 후 이분탐색을 사용하여 최적화 하는 방법이 떠올랐다.

 

 

정답코드

#include <bits/stdc++.h>
using namespace std;
typedef pair<int, int> pii;
typedef long long ll;
#define endl "\n"
#define MAX 1e9

struct coordinate {
    int x;
    int y;
    int r;
};

struct horse {
    int x;
    int y;
    int dir;
};

// int dx[] = {-1, 1, 0, 0, 1, -1, -1, 1};
// int dy[] = {0, 0, -1, 1, -1, 1, -1, 1};
int dx[] = {0, 0, 1, -1};
int dy[] = {1, -1, 0, 0};
int N;
int A[4001];
int B[4001];
int C[4001];
int D[4001];


void solve() {
    vector<ll> AB;
    vector<ll> CD;

    for(int i = 0; i < N; i++) {
        for(int j = 0; j < N; j++) {
            ll tmp = A[i] + B[j];
            ll tmps = C[i] + D[j];
            AB.push_back(tmp);
            CD.push_back(tmps);
        }
    }

    ll cnt = 0;
    sort(AB.begin(), AB.end());
    sort(CD.begin(), CD.end());

    for(int i = 0; i < AB.size(); i++) {
        ll tmp = -AB[i];

        ll low = lower_bound(CD.begin(), CD.end(), tmp) - CD.begin();
        ll high = upper_bound(CD.begin(), CD.end(), tmp) - CD.begin();

        if(tmp == CD[low]) cnt += (high - low);
    }

    cout << cnt;
}

void input() {
    cin >> N;

    for(int i = 0; i < N; i++) {
        for(int j = 0; j < 4; j++) {
            if(j == 0) {
                cin >> A[i];
            }
            else if(j == 1) cin >> B[i];
            else if(j == 2) cin >> C[i];
            else cin >> D[i];
        }
    }
}

int main() {
    ios::sync_with_stdio(0);
    cin.tie(0), cout.tie(0);

    input();
    solve();
}