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

https://www.acmicpc.net/problem/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을 출력한다.

 

문제 풀이

---

upper_bound와 lower_bound를 사용하는 이분 탐색 문제.

A와 B 각각 모든 가능한 부분합을 구한 후 벡터(v1, v2)에 저장해놓는다.

upper_bound는 찾는 값보다 큰 값을 가지는 것 중 가장 빠른 iterator를 반환하고 lower_bound는 찾는 값과 같거나 큰 값 중 가장 빠른 index의 iterator를 반환하므로 upper_bound 값에서 lower_bound 값을 빼면 찾는 값이 벡터(혹은 배열)에 몇 개 존재하는지 찾을 수 있다.

v1의 처음부터 끝까지 탐색하면서 v2에 T - v1(i)값이 몇개 있는지 upper_bound- lower_bound로 찾아주면 된다.

 

아래는 코드.

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;
int N, M;
long long answer;

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

	int T;
	cin >> T;
	cin >> N;
	int* arrA = new int[N];
	int* partSumA = new int[N + 1];
	partSumA[0] = 0;
	vector<int> v1;
	vector<int> v2;
	for (int i = 0; i < N; i++)
	{
		cin >> arrA[i];
		partSumA[i + 1] = partSumA[i] + arrA[i];
	}
	for (int i = 1; i < N+1; i++)
	{
		for (int j = 0; j < i; j++)
		{
			v1.push_back(partSumA[i] - partSumA[j]);
		}
	}
	cin >> M;
	int* arrB = new int[M];
	int* partSumB = new int[M + 1];
	partSumB[0] = 0;
	for (int i = 0; i < M; i++)
	{
		cin >> arrB[i];
		partSumB[i + 1] = partSumB[i] + arrB[i];
	}
	for (int i = 1; i < M + 1; i++)
	{
		for (int j = 0; j < i; j++)
		{
			v2.push_back(partSumB[i] - partSumB[j]);
		}
	}
	sort(v2.begin(), v2.end());
	for (int i = 0; i < v1.size(); i++)
	{
		answer += upper_bound(v2.begin(), v2.end(), T - v1.at(i)) - lower_bound(v2.begin(), v2.end(), T - v1.at(i));
	}
	cout << answer << "\n";
	return 0;
}