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

Baekjoon 2143

문제

숫자 T가 주어지고 배열 A와 B가 주어진다. 여기서 부 배열이라는 개념이 있는데 A[i], A[i+1], ..., A[j] 와 같이 연속된 배열을 부배열이라고 한다. A[13]까지라고 한다면 A[7], A[8], A[9], ... A[13] 도 가능하고 A[2], A[3]도 부 배열이라고 할 수 있다.

 

문제는 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]

아이디어

배열의 개수는 A, B 2개로 제한되어 있다. 여기서 방법은 다음과 같다.
A에서 구할 수 있는 부 배열의 합의 모든 경우의 수의 빈도를 해시 맵에 저장한다. {"sum": "frequency"}처럼 저장하면 된다. 예를 들어, 부 배열의 합이 13으로 한 번 있으면 {"13": "1"}, 부 배열의 합이 15로 5번 있으면 {"15", "5"}로 저장한다.

 

그 다음으로 B의 부 배열의 모든 경우의 수를 구하는 데 A를 해시 맵에 저장한 것처럼 하지 않고 A의 해시맵에서 합의 수를 T에서 빼준 값을 판별한다. 즉, B의 부배열의 합이 T-(A 부 배열의 합)이 되도록한다. 경우의 수가 몇 번인지 계산 하면 되기 때문에 A의 해시 맵에 빈도 수만 저장해주었다. 이제 B의 부배열의 합이 T-(A 부배 열의 합)이 된다면 해당 A 해시 맵의 빈도수를 count에 더해주면 된다.

CODE

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

typedef long long ll;

void init() {
    cin.tie(NULL);
    cout.tie(NULL);
    ios::sync_with_stdio(false);
}

int main() {
    init();
    int T;
    cin >> T;
    int n;
    cin >> n;
    vector<int> A(n);
    for (int i = 0; i < n; i++) {
        cin >> A[i];
    }
    int m;
    cin >> m;
    vector<int> B(m);
    for (int i = 0; i < m; i++) {
        cin >> B[i];
    }
    // 부 배열의 합의 빈도를 저장하는 해시 테이블
    unordered_map<int, int> sumA;
    for (int i = 0; i < n; i++) {
        int current_sum = 0;
        for (int j = i; j < n; j++) {
            current_sum += A[j];
            sumA[current_sum]++;
        }
    }

    ll count = 0;
    // B의 모든 부 배열의 합을 계산하고 sumA에서 T와의 차이를 찾는다.
    for (int i = 0; i < m; i++) {
        int current_sum = 0;
        for (int j = i; j < m; j++) {
            current_sum += B[j];
            int required_sum = T - current_sum;
            if (sumA.find(required_sum) != sumA.end()) {
                // A의 빈도를 더해주면 A에서 경우의 수를 굳이 세지 않아도 된다.
                count += sumA[required_sum];
            }
        }
    }

    cout << count << endl;

    return 0;
}