[백준] 2143번 : 두 배열의 합 (C++)

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을 출력한다.

 

 

 

 

 

풀이)

 

// 2143. 두 배열의 합
#include <iostream>
#include <vector>
#include <algorithm>
#include <map>
 
using namespace std;
 
int main()
{
    ios_base::sync_with_stdio(false);
    cin.tie(0);
 
    long long t;
    cin >> t;
 
    int n, m;
    cin >> n;
 
    // 원본배열은 누적합을 만들 때 말고는 필요없으므로 
    // 바로 누적합만 만들어준다.
    vector<long long> prefix1(n + 1, 0);
    for (int i = 1; i < n + 1; i++)
    {
        cin >> prefix1[i];
        prefix1[i] += prefix1[i - 1];
    }
 
    // 해당 누적합 배열을 2중 for문을 통해
    // 만들 수 있는 모든 합을 만들어 map에 넣어준다.
    // map을 사용하는 이유는 중복되는 값이 있을 수 있기 때문에
    // 효율과 메모리를 위해 사용한다.
    map<long long, long long> map1;
    for (int i = 1; i < n + 1; i++)
    {
        for (int j = i; j < n + 1; j++)
        {
            map1[prefix1[j] - prefix1[i - 1]]++;
        }
    }
 
    cin >> m;
    vector<long long> prefix2(m + 1, 0);
    for (int i = 1; i < m + 1; i++)
    {
        cin >> prefix2[i];
        prefix2[i] += prefix2[i - 1];
    }
 
    map<long long, long long> map2;
    for (int i = 1; i < m + 1; i++)
    {
        for (int j = i; j < m + 1; j++)
        {
            map2[prefix2[j] - prefix2[i - 1]]++;
        }
    }
 
 
    // map을 vector로 변환해 투포인터를 하면 편하지만,
    // 문제의 메모리 제한이 작기 때문에 iterator를 이용해 
    // 투포인터를 진행한다.
    auto s = map1.begin();
    auto e = map2.end();
    e--;
 
    long long ans = 0;
    while (!(s == map1.end() && e == map2.begin()))
    {
        auto a = s;
        if (s == map1.end()) a--;
 
        long long _sum = a->first + e->first;
        if (_sum == t) ans += ((a->second) * (e->second));
 
        if (s == map1.end()) e--;
        else if (e == map2.begin()) s++;
        else if (_sum < t) s++;
        else if (_sum >= t) e--;
    }
 
    cout << ans;
 
    return 0;
}
 

---

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

2143번: 두 배열의 합