BOJ 2143 - 두 배열의 합

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

2143번: 두 배열의 합

본 문제는 $O(nm)$대신, $O(nlogm)$으로 접근할 수 있는 방법이 있습니다.

($n$, $m$은 $A_p$, $B_p$의 길이)

 

먼저, 배열 $A$, $B$에 대해서 가능한 모든 부분합을 $A_p$, $B_p$에 저장하는 과정을 각각 $O(N^2)$, $O(M^2)$으로 구해줍니다. 

그 다음, $A_p$, $B_p$를 sort하여 lower_bound, upper_bound를 사용할 수 있게 세팅합시다. ($O(nlogn), O(mlogm)$)

 

그 다음, $A_p$의 모든 원소에 대해($O(n)$) $T-{A_p}_i$의 값이 $B_p$에 몇 개가 존재하는 지를 upper_bound-lower_bound의 값으로 구해 cnt에 더해줍니다. ($O(nlogm)$)

 

ll arr[1005];
ll brr[1005];
ll asum[1005];
ll bsum[1005];

vector<ll> a_exist;
vector<ll> b_exist;

int main() {
    fastio();
    ll T;
    cin >> T;
    int N;
    cin >> N;
    for(int i=0; i<N; i++) {
        cin >> arr[i];
    }
    int M;
    cin >> M;
    for(int i=0; i<M; i++) {
        cin >> brr[i];
    }

    asum[0] = arr[0];
    for(int i=1; i<N; i++) {
        asum[i] = asum[i-1] + arr[i];
    }
    bsum[0] = brr[0];
    for(int i=1; i<M; i++) {
        bsum[i] = bsum[i-1] + brr[i];
    }

    for(int i=0; i<N; i++) {
        for(int j=i; j<N; j++) {
            if(i==0) {
                a_exist.push_back(asum[j]);
            } else {
                a_exist.push_back(asum[j]-asum[i-1]);
            }
        }
    }
    for(int i=0; i<M; i++) {
        for(int j=i; j<M; j++) {
            if(i==0) {
                b_exist.push_back(bsum[j]);
            } else {
                b_exist.push_back(bsum[j]-bsum[i-1]);
            }
        }
    }
    sort(all(a_exist));
    sort(all(b_exist));
    ll cnt=0;
    for(int i=0; i<a_exist.size(); i++) {
        ll X = T - a_exist[i];
        cnt += upper_bound(all(b_exist), X) - lower_bound(all(b_exist), X);
    }
    cout << cnt;
}