문제 간략 설명

각 원소가 정수인 두 배열 A[1], …, A[n]과 B[1], …, B[m]이 주어졌을 때, A의 부배열의 합에 B의 부배열의 합을 더해서 T가 되는 모든 부배열 쌍의 개수를 구하는 프로그램을 작성하시오.

풀이 방법

A, B 배열의 부분합을 map_a, map_b에 저장한다. (key: 부분합, value: 부분합 개수)

map_a를 순회하면서 map_b[T - iter.first]에 접근하여 (합해서 T가 되게) 곱해서 res를 구하고 출력한다.

코드

#include <iostream>
#include <vector>
#include <map>

using namespace std;

typedef long long ll;

int T, N, M;
map<int, int> map_a;
map<int, int> map_b;
vector<int> vec;

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

	cin >> T;

	cin >> N;

	vec.resize(N);

	for (int i = 0; i < N; i++)
		cin >> vec[i];

	for (int i = 0; i < N; i++) {
		int sum = 0;
		for (int j = i; j < N; j++) {
			sum += vec[j];
			map_a[sum]++;
		}
	}

	cin >> N;

	vec.clear();
	vec.resize(N);

	for (int i = 0; i < N; i++)
		cin >> vec[i];

	for (int i = 0; i < N; i++) {
		int sum = 0;
		for (int j = i; j < N; j++) {
			sum += vec[j];
			map_b[sum]++;
		}
	}
	
	ll res = 0;

	for (auto iter: map_a)
		res += (ll)iter.second * map_b[T - iter.first]; 

	cout << res << '\\n';

	return 0;
}

Key Points