백준 2143 두 배열의 합 (C++)

2143번: 두 배열의 합

누적 합, 이분 탐색을 이용한 문제이다. 먼저 A와 B의 모든 누적 합들을 구해준다. 그리고 B의 누적합을 오름차순으로 정렬을 한 후 A의 누적 합 크기만큼 반복문을 돌면서 T와의 차이와 같은 값을 B 누적 합에서 찾아 시작과 끝 인덱스를 구해준다. 시작과 끝 인덱스 차이를 더해주며 이를 반복한 후 출력해주었다. 굉장히 어려웠던 문제였다. 이분 탐색을 이용한다는 점은 생각해냈었는데 누적 합을 이용하는 점까진 생각을 하지 못해 많이 고생을 했다. 그리고 이분 탐색을 lower_bound와 upper_bound를 이용한다는 점도 알지 못했던 정보였다. 많은 것들을 알 수 있었던 문제였다. 누적 합과 관련된 문제들을 좀 더 풀어봐야 겠다.

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

using namespace std;

int T, n, m;
long long result = 0;
int A[1000];
int B[1000];
vector<int> va, vb;

void solution() {
    for (int i = 0; i < n; i++) {
        int sum = A[i];
        va.push_back(sum);
        for (int j = i + 1; j < n; j++) {
            sum += A[j];
            va.push_back(sum);
        }
    }

    for (int i = 0; i < m; i++) {
        int sum = B[i];
        vb.push_back(sum);
        for (int j = i + 1; j < m; j++) {
            sum += B[j];
            vb.push_back(sum);
        }
    }

    sort(vb.begin(), vb.end());

    for (int i = 0; i < va.size(); i++) {
        int dif = T - va[i];
        auto lo = lower_bound(vb.begin(), vb.end(), dif);
        auto up = upper_bound(vb.begin(), vb.end(), dif);

        result += (up - lo);
    }

    cout << result;
}

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

    cin >> T;
    cin >> n;

    for (int i = 0; i < n; i++) {
        cin >> A[i];
    }

    cin >> m;

    for (int i = 0; i < m; i++) {
        cin >> B[i];
    }

    solution();

    return 0;
}