백준 2143번 두 배열의 합 C++

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

 

우선 A와 B의 모든 구간에서의 누적합들을 저장하는 벡터를 만든다.

 

A의 누적합과 B의 누적합을 더했을때 t가 되는 경우를 찾으면 되는데, 이분 탐색을 위해 둘중 하나를 정렬한다. 

나는 B를 정렬했다.

 

A의 모든 원소를 탐색하면서 B의 원소의 값이 t - A[i]인 경우를 찾으면 되는데, 이때 lower_bound와 upper_bound를 사용했다. lower_bound와 upper_bound는 이분 탐색을 해주는 함수인데, 정렬이 되어있어야 사용이 가능하다. 그래서 앞에서 B를 정렬했다.

 

lower_bound는 찾고싶은 수보다 같거나 큰 인덱스를 반환하고

upper_bound는 찾고싶은 수보다 큰 인덱스를 반환한다.

 

ans에 (upper_bound의 반환값 - lower_bound의 반환값) 을  더해주면 된다.

그 이유는 만약 두 반환값이 같다면 찾고싶은 수와 같은 경우는 없어서 0이 되고

반환값이 다르다면 찾고싶은 수가 upper_bound의 반환값 - lower_bound의 반환값 만큼 있기 때문이다.

 

#include <iostream>
#include <string>
#include <vector>
#include <math.h>
#include <algorithm>
#include <utility>
#include <stack>
#include <queue>
#include <math.h>
#include <set>
#include <map>
#include <list>
#include <unordered_map>
#include <unordered_set>
#include <iomanip>
#include <limits.h>

using namespace std;
using int64 = long long;

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    cout.tie(NULL);

    int t, n, m;
    cin >> t;
    cin >> n;
    vector<int> A(n);
    for (int i = 0; i < n; i++) {
        cin >> A[i];
    }
    cin >> m;
    vector<int> B(m);
    for (int i = 0; i < m; i++) {
        cin >> B[i];
    }
    vector<int> ASum;
    vector<int> BSum;
    for (int i = 0; i < n; i++) {
        int sum = A[i];
        ASum.push_back(sum);
        for (int j = i + 1; j < n; j++) {
            sum += A[j];
            ASum.push_back(sum);
        }
    }
    for (int i = 0; i < m; i++) {
        int sum = B[i];
        BSum.push_back(sum);
        for (int j = i + 1; j < m; j++) {
            sum += B[j];
            BSum.push_back(sum);
        }
    }
    sort(BSum.begin(), BSum.end());

    int64 ans = 0;
    for (auto a : ASum) {
        int target = t - a;
        auto lb = lower_bound(BSum.begin(), BSum.end(), target);
        auto ub = upper_bound(BSum.begin(), BSum.end(), target);
        ans += ub - lb;
    }
    cout << ans;
}