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

C++ 코드

#include <iostream>
#include <vector>
#include <algorithm>

using namespace std;

int main() {
	long long t;
	cin >> t;
	int n, m;
	cin >> n;
	/////// 입력 //////
	vector<int> arr1(n);
	for (int i = 0; i < n; i++) {
		cin >> arr1[i];
	}
	cin >> m;
	vector<int> arr2(m);
	for (int i = 0; i < m; i++) {
		cin >> arr2[i];
	}
	////// 부 집합 생성 /////
	vector<int> first;
	for (int i = 0; i < n; i++) {
		int sum = 0;
		for (int k = i; k < n; k++) {
			sum += arr1[k];
			first.push_back(sum);
		}
	}
	vector<int> second;
	for (int i = 0; i < m; i++) {
		int sum = 0;
		for (int k = i; k < m; k++) {
			sum += arr2[k];		
			second.push_back(sum);
		}
	}

	sort(first.begin(), first.end());
	sort(second.begin(), second.end());
	long long ans = 0;
	for (int i = 0; i < first.size(); i++) {
		auto p = equal_range(second.begin(), second.end(), t - first[i]);
		// 임의의 값이랑 같은 구간의 시작과 끝을 리턴
		ans += p.second - p.first;
	}

	cout << ans << '\n';



	return 0;
}