BOJ - 2143 두 배열의 합

두 배열이 주어지고 각 배열의 부배열의 합끼리 더한 값이 T가 되는 경우의 수를 구하는 문제이다.

문제풀이 전략

문제에서 주어진 부배열이란 i부터 j까지 연속된 부분배열을 의미한다.
즉, 부배열의 합은 A[i] ~ A[j] 까지 합친 값을 의미한다.

우선 배열 A,B에 대하여 부배열이 될 수 있는 모든 경우의 부배열 합을 구한다.
예를들어 배열 A가 1, 2, 3, 4, 5 라면 1, 1+2, 1+2+3, ..., 4, 4+5, 5 를 구한다.

부배열의 합을 모두 구한 것을 투포인터를 사용하거나 이분탐색을 사용하여 문제를 해결할 수 있다.
여기서는 이분탐색을 사용하였다.

두 부배열의 합 중 하나를 오름차순으로 정렬한다.

그리고 나서 T에서 정렬이 안된 부배열의 합을 뺀 값을 정렬된 부배열의 합에서 찾으면 된다.

이 과정에서 upper_bound와 lower_bound를 이용하였다.
(참고 https://chanhuiseok.github.io/posts/algo-55/)

upper_bound는 초과되는 인덱스를, lower_bound는 이상인 인덱스를 찾으므로 두 인덱스의 차를 구하면 일치하는 값의 갯수가 된다.

코드

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

using namespace std;

int t, n, m;
vector<int> A;
vector<int> sA;
vector<int> B;
vector<int> sB;
int main(){
    cin >> t;
    cin >> n;
    for(int i=0;i<n;i++){
        int a;
        cin >> a;
        A.push_back(a);
    }
    for(int i=0;i<n;i++){
        int sum = 0;
        for(int j=i;j<n;j++){
            sum += A[j];
            sA.push_back(sum);
        }
    }
    cin >> m;
    for(int i=0;i<m;i++){
        int b;
        cin >> b;
        B.push_back(b);
    }
    for(int i=0;i<m;i++){
        int sum = 0;
        for(int j=i;j<m;j++){
            sum += B[j];
            sB.push_back(sum);
        }
    }

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

    long long ans = 0;
    for(int i=0;i<sA.size();i++){
        int f = t - sA[i];
        auto up = upper_bound(sB.begin(), sB.end(), f);
        auto lp = lower_bound(sB.begin(), sB.end(), f);
        ans += up - lp;
    }
    cout << ans << "\n";
}

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