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

문제

한 배열 A[1], A[2], …, A[n]에 대해서, 부 배열은 A[i], A[i+1], …, A[j-1], A[j] (단, 1 ≤ i ≤ j ≤ n)을 말한다. 이러한 부 배열의 합은 A[i]+…+A[j]를 의미한다. 각 원소가 정수인 두 배열 A[1], …, A[n]과 B[1], …, B[m]이 주어졌을 때, A의 부 배열의 합에 B의 부 배열의 합을 더해서 T가 되는 모든 부 배열 쌍의 개수를 구하는 프로그램을 작성하시오.

예를 들어 A = {1, 3, 1, 2}, B = {1, 3, 2}, T=5인 경우, 부 배열 쌍의 개수는 다음의 7가지 경우가 있다.

T(=5) = A[1] + B[1] + B[2]
= A[1] + A[2] + B[1]
= A[2] + B[3]
= A[2] + A[3] + B[1]
= A[3] + B[1] + B[2]
= A[3] + A[4] + B[3]
= A[4] + B[2]

입력

첫째 줄에 T(-1,000,000,000 ≤ T ≤ 1,000,000,000)가 주어진다. 다음 줄에는 n(1 ≤ n ≤ 1,000)이 주어지고, 그 다음 줄에 n개의 정수로 A[1], …, A[n]이 주어진다. 다음 줄에는 m(1 ≤ m ≤ 1,000)이 주어지고, 그 다음 줄에 m개의 정수로 B[1], …, B[m]이 주어진다. 각각의 배열 원소는 절댓값이 1,000,000을 넘지 않는 정수이다.

출력

첫째 줄에 답을 출력한다. 가능한 경우가 한 가지도 없을 경우에는 0을 출력한다.

예제 입출력

• 예제 입력 1

5
4
1 3 1 2
3
1 3 2

• 예제 출력 1

7

Solution

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;

int main(){
    ios::sync_with_stdio(0);
    cin.tie(0);
    long long T;
    int N, M;
    vector<long long> A, B;
    vector<long long> ASum, BSum;
    long long tmp, sum, result =0;
    cin >> T;
    cin >> N;
    for(int i = 0; i < N; i++){
        cin >> tmp;
        A.push_back(tmp);
    }
    cin >> M;
    for(int i = 0; i < M; i++){
        cin >> tmp;
        B.push_back(tmp);
    }
    for(int i = 0; i < N; i++){
        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++){
        sum = B[i];
        BSum.push_back(sum);
        for(int j = i + 1; j < M; j++){
            sum += B[j];
            BSum.push_back(sum);
        }
    }
    sort(ASum.begin(), ASum.end());
    sort(BSum.begin(), BSum.end());
    for (int i = 0; i < ASum.size(); i++){
        int low = lower_bound(BSum.begin(), BSum.end(), T - ASum[i]) - BSum.begin();
        int high = upper_bound(BSum.begin(), BSum.end(), T - ASum[i]) - BSum.begin();
        result += high - low;
    }
    cout << result << "\n";
}

누적합을 쓰는 이유는 다양하게 있겠지만 가장 중요한 이유는 중복되는 계산을 하지 않기 위해서라고 생각한다.

데이터의 범위를 생각했을 때 long long 데이터들로 처리를 해야 한다는 것을 알 수 있다.

 for(int i = 0; i < N; i++){
    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++){
    sum = B[i];
    BSum.push_back(sum);
    for(int j = i + 1; j < M; j++){
        sum += B[j];
        BSum.push_back(sum);
    }
}
sort(ASum.begin(), ASum.end());
sort(BSum.begin(), BSum.end());
for (int i = 0; i < ASum.size(); i++){
    int low = lower_bound(BSum.begin(), BSum.end(), T - ASum[i]) - BSum.begin();
    int high = upper_bound(BSum.begin(), BSum.end(), T - ASum[i]) - BSum.begin();
    result += high - low;
}
cout << result << "\n";

특정 값 기준으로 그 값 이후의 값들을 모두 계산하는 모든 경우를 누적 합 처리한다. 또한 이 값들을 모두 정렬 해 줌으로써 탐색을 빠르게 할 수 있도록 처리한다.

A를 기준으로 탐색하면서 B의 누적합데이터에서 A의 i번째 누적합 데이터와 더했을 때 T가 되는 데이터 최대 최소값을 계산한다. 중복되는 데이터가 충분히 있을 수 있기 때문이다. 이 값들의 위치 데이터를 계산해서 계속해서 누적 해 나가면 정답이 나온다.

굳이 값들을 토대로 탐색을 할 필요는 없다. 누적합을 활용할 수 있는 좋은 문제였다.