[백준 2143] 두 배열의 합 C++

2143번: 두 배열의 합

이분 탐색

---

1. 문제해결 아이디어

문제에서 정의된 부 배열의 최대 크기는 N(N+1)/2 가 됨을 알 수 있었다. (1~N까지의 합)

N, M이 최대 1000이므로 부 배열의 크기는 최대 500,000 이다.

따라서 배열에 담은 후 nlogn 연산이 충분히 가능함을 알 수 있었다.

 

이분 탐색을 수행하기 위해 subB배열을 정렬한다.

T 에서 subA의 원소를 뺀 값이 subB에 있는지 알아보기 위해 lower_bound, upper_bound를 수행해서 값을 찾는다.

 

총 시간복잡도는  n^2 + n^2 + nlogn + n^2 * logn 이다. (subA구하기, subB구하기, subB정렬, subA를 돌면서 subB이분탐색)

---

2. 코드

<+12> : 음수가 들어올 줄 모르고 걸어둔 bound 조건이다.. 덕분에 많이 틀렸다.

<+18> ~ <+22> : 처음엔 upper_bound에서 구한 itr에서 하나씩 내려가며 tw값과 같은지 비교하고, 같다면 cnt++를 수행하는 코드를 짰었다. 그런데 그럴 필요없이 upper_bound - lower_bound 하면 그 사이의 같은 값들의 수를 바로 알 수 있었다.

<+61> : subA는 이분탐색을 수행하지 않으므로 정렬 할 필요가 없다.

 

*주의사항 : long long 타입 필수 

 

#include <bits/stdc++.h>
 
using namespace std;
 
int T, N, M;
long long A[1001], B[1001], PA[1001], PB[1001];
vector<long long> subA, subB;
 
void solve(){
    long long tw, cnt = 0;
    for(auto w : subA){
//        if(w >= T) break;
        tw = T - w;
        
        auto itrl = lower_bound(subB.begin(), subB.end(), tw);
        auto itru = upper_bound(subB.begin(), subB.end(), tw);
        cnt += (itru - itrl);
//        itr--;
//        for(; itr >= subB.begin(); itr--){
//            if(*itr != tw) break;
//            cnt++;
//        }
    }
    
    cout << cnt;
}
 
void init(){
    cin >>T;
    cin >> N;
    for(int i = 1; i <= N; i++){
        cin >> A[i];
    }
    cin >> M;
    for(int i = 1; i <= M; i++){
        cin >> B[i];
    }
    
    // 누적합 배열 구하기
    PA[0] = 0;
    PB[0] = 0;
    for(int i = 1; i <= N; i++){
        PA[i] = PA[i-1] + A[i];
    }
    for(int i = 1; i <= M; i++){
        PB[i] = PB[i-1] + B[i];
    }
    
    // 모든 부 배열의 합 구하기
    for(int i = 1; i <= N; i++){
        for(int j = i; j <= N; j++){
            subA.push_back(PA[j] - PA[j-i]);
        }
    }
    for(int i = 1; i <= M; i++){
        for(int j = i; j <= M; j++){
            subB.push_back(PB[j] - PB[j-i]);
        }
    }
    
//    sort(subA.begin(), subA.end());
    sort(subB.begin(), subB.end());
}
 
int main(){
    init();
    solve();
    
    return 0;
}