[BOJ] 2143 두 배열의 합 - C++

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

2143번: 두 배열의 합

아이디어

두 수열 $A_n$, $B_m$의 연속한 원소들의 합을 $T$가 되도록 만들면 된다. 누적합을 통해 연속한 원소의 합을 빠르게 구할 수 있고 이분탐색을 통해 가능한 경우의 수를 빠르게 셀 수 있다. 부분합은 $1$부터 $i$까지의 합이고 이번 문제에서는 $i$부터 $j$까지의 합이 필요하므로 미리 나올 수 있는 모든 합의 경우를 구해놓으면 된다.

 

참고로 정렬되어 있는 배열에서 특정 값의 개수를 새려면 upper_bound에서 lower_bound를 빼는 방법이 있지만 equal_range를 통해 pair로 시작점과 끝점을 가져오는 방법도 있다. upper_bound와 lower_bound를 한 번에 쓴 것과 같다.

 

또, 숫자의 범위가 int를 벗어나는 것에 주의해야 한다.

코드

#include <iostream>
#include <algorithm>
#include <vector>

using namespace std;

typedef long long ll;

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

    int t;
    cin >> t;
    int n;
    cin >> n;
    ll a[n];
    for (int i = 0; i < n; i++) {
        if (i == 0) {
            cin >> a[i];
            continue;
        }
        int num;
        cin >> num;
        a[i] = a[i - 1] + num;
    }

    int m;
    cin >> m;
    ll b[m];
    for (int i = 0; i < m; i++) {
        if (i == 0) {
            cin >> b[i];
            continue;
        }
        int num;
        cin >> num;
        b[i] = b[i - 1] + num;
    }

    vector<ll> a_sum;
    vector<ll> b_sum;

    for (int i = 0; i < n; i++) {
        for (int j = i; j < n; j++) {
            a_sum.push_back(a[j] - (i == 0 ? 0 : a[i - 1]));
        }
    }
    sort(a_sum.begin(), a_sum.end());

    for (int i = 0; i < m; i++) {
        for (int j = i; j < m; j++) {
            b_sum.push_back(b[j] - (i == 0 ? 0 : b[i - 1]));
        }
    }
    sort(b_sum.begin(), b_sum.end());

    ll ans = 0;
    for (long long i : a_sum) {
        ll target = t - i;
        auto range = equal_range(b_sum.begin(), b_sum.end(), target);
        ans += range.second - range.first;
    }

    cout << ans;
}