2143. 두 배열의 합

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

2143번: 두 배열의 합

문제 설명

이 문제는 두 배열 \( A, B \)의 연속하는 부분배열의 합이 T가 되도록하는 경우의 수를 구하는 문제이다.

 

\( N <= 1000 \) 인 문제로 각 배열의 모든 누적합 요소를 가지고 있어도 \( \frac{N(N+1)}{2} \)개 이므로 최대 500,500개의 요소를 각 배열에 가지게 된다.

 

먼저 누적합 배열 \( pxA, pxB \) 을 각각 구성해야한다.

여기서 누적합 배열에는 우리가 이전에 알고 있던 누적합 요소 뿐만 아니라 모든 누적합 요소를 가지고 있어야한다. (두 배열 A, B의 요소가 자연수가 아닌 정수이기 때문)

 

모든 누적합 요소를 갖춘 배열 2개로 이분탐색을 진행한다.

\( T \)에서 \( pxA[i] \)을 뺀 값을 \( pxB \) 배열에서 찾으면 된다.

같은 값이 여러개 있을 수 있으므로 upper_bound값과 lower_bound값을 찾아 둘을 뺀 값을 모두 더한 값이 이 문제의 답이 된다.

 

이분탐색은 C++ STL의 lower_bound와 upper_bound를 사용했다.

#include <bits/stdc++.h>
 
using namespace std;
 
#define swap(a,b) (a)^=(b)^=(a)^=(b)
#define endl '\n'
#define int long long
 
using pii = pair<int, int>;
 
const double EPS = 1e-9;
const int dx[] = {1, -1, 0, 0};
const int dy[] = {0, 0, 1, -1};
 
 
 
int32_t main()
{
    cin.sync_with_stdio(NULL);
    cin.tie(NULL);
    cout.tie(NULL);
 
    int T, n, m;
    vector<int> A, B;
    cin >> T;
    cin >> n;
    A.resize(n, 0);
    for(int i=0;i<n;i++)
        cin >> A[i];
    cin >> m;
    B.resize(m, 0);
    for(int i=0;i<m;i++)
        cin >> B[i];
    vector<int> pxA(n, 0), pxB(m, 0);
    pxA[0] = A[0];
    pxB[0] = B[0];
    for(int i=1;i<n;i++)
        pxA[i] = pxA[i-1] + A[i];
    for(int i=1;i<m;i++)
        pxB[i] = pxB[i-1] + B[i];
    for(int i=0;i<n;i++)
        for(int j=i+1;j<n;j++)
            pxA.push_back(pxA[j]-pxA[i]);
    for(int i=0;i<m;i++)
        for(int j=i+1;j<m;j++)
            pxB.push_back(pxB[j]-pxB[i]);
    sort(pxA.begin(), pxA.end());
    sort(pxB.begin(), pxB.end());
 
    int ans = 0;
    for(int& pxak : pxA)
    {
        auto lit = lower_bound(pxB.begin(), pxB.end(), T-pxak);
        auto uit = upper_bound(pxB.begin(), pxB.end(), T-pxak);
        ans += uit-lit;
    }
    cout << ans;
 
    return 0;
}