[Baekjoon/C++] 2143번 - 두 배열의 합

[백준] Baekjoon Online Judge

문제로 이동

 

문제

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

 

---

풀이

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

int T, n, m, num[1001];
vector<long long> aSum;
unordered_map<long long, int> bSum;

void makeASum();
void makeBSum();

void makeASum() {
    for (int i = 0; i < n; i++) {
        int sum = 0;

        for (int j = i; j < n; j++) {
            sum += num[j];
            aSum.push_back(sum);
        }
    }
}

void makeBSum() {
    for (int i = 0; i < m; i++) {
        int sum = 0;

        for (int j = i; j < m; j++) {
            sum += num[j];
            bSum[sum]++;
        }
    }
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(0);
    cout.tie(0);

    cin >> T >> n;
    for (int i = 0; i < n; i++) cin >> num[i];

    makeASum();

    cin >> m;
    for (int i = 0; i < m; i++) cin >> num[i];

    makeBSum();

    long long result = 0;

    for (int i = 0; i < aSum.size(); i++) {
        if (bSum.find(T - aSum[i]) != bSum.end())
            result += bSum[T - aSum[i]];
    }

    cout << result << '\n';

    return 0;
}

 

void makeASum() {
    for (int i = 0; i < n; i++) {
        int sum = 0;

        for (int j = i; j < n; j++) {
            sum += num[j];
            aSum.push_back(sum);
        }
    }
}

void makeBSum() {
    for (int i = 0; i < m; i++) {
        int sum = 0;

        for (int j = i; j < m; j++) {
            sum += num[j];
            bSum[sum]++;
        }
    }
}

 각 배열에서 나올 수 있는 부 배열의 합을 미리 aSum과 bSum에 계산하도록 했다. bSum의 값이 추가적인 코드 없이 바로 나왔으면 좋겠다고 생각했고, 정렬이 필요 없었기에 unordered_map을 사용하였다.

 

    long long result = 0;

    for (int i = 0; i < aSum.size(); i++) {
        if (bSum.find(T - aSum[i]) != bSum.end())
            result += bSum[T - aSum[i]];
    }

 두 배열의 부분합으로 T를 만들 수 있는 쌍의 개수를 구하는 부분의 코드이다. map은 존재하지 않는 키를 사용하여 호출하면 이를 map에 새로 추가한다. 그렇기에 메모리 제한을 지키기 위해 if문으로 원하는 값이 map에 존재하는지 확인하도록 했다.