[BOJ] 2143 : 두 배열의 합

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

풀이

누적합과 맵을 활용하여 풀 수 있었다.
이전에 풀어보았던 BOJ - 2015 수들의 합 문제와 유사한 방법으로 풀었고,
문제에 주어진 조건 T, n, m 의 범위를 잘 고려하지 못하여 68% 에서 몇 번의 시행착오를 겪었다.

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

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

	int t;
	cin >> t;
	int n, m;
	cin >> n;

	map<long long, long long> mp;

	vector<long long> a(n + 1);

	for (int i = 1; i <= n; ++i)
	{
		cin >> a[i];
		a[i] += a[i - 1];

		for (int j = 0; j < i; ++j)
		{
			long long temp = a[i] - a[j];
			mp[t - temp]++;
		}
	}

	cin >> m;
	vector<long long> b(m + 1);

	long long answer = 0;

	for (int i = 1; i <= m; ++i)
	{
		cin >> b[i];
		b[i] += b[i - 1];

		for (int j = 0; j < i; ++j)
		{
			long long temp = b[i] - b[j];
			answer += mp[temp];
		}
	}

	cout << answer;
	return 0;
}