[C/C++] 백준 - 2143번 : 두 배열의 합

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

 

처음에 부 배열의 정의를 놓쳐서 헤맸던 문제이다. 부 배열은 연속된 수들의 집합이다. 

이것만 알면 해결할 수 있다. 처음에 입력했던 각 배열에 누적합을 모두 추가해준 뒤 정렬해주고, (2개 모두)

하나의 배열 원소 길이만큼 T에서 A[i]를 뺀 값이 B배열에 몇 개가 존재하는지 lower_bound와 upper_bound를 통해 구한뒤 개수를 더해주면 된다.

#include<iostream>
#include<cmath>
#include<list>
#include<stack>
#include<tuple>
#include<map>
#include<algorithm>
#include<vector>
#include<deque>
#include<queue>
#define CUNLINK ios::sync_with_stdio(false), cin.tie(0), cout.tie(0)
#define endl '\n'
#define MAX_SIZE 1003
#define INF 987654321
using namespace std;
typedef long long ll;

int main() {
	CUNLINK;
	vector<ll> A, B;
	ll n_A, n_B, T;
	
	cin >> T;
	cin >> n_A; 
	for (int i = 0; i < n_A; i++) {
		int x; cin >> x;
		A.push_back(x);
	}
	cin >> n_B;
	for (int i = 0; i < n_B; i++) {
		int x; cin >> x;
		B.push_back(x);
	}
		
	for (int i = 0; i < n_A; i++) {
		ll Sum = A[i]; 
		for (int j = i + 1; j < n_A; j++) {
			Sum += A[j];
			A.push_back(Sum);
		}
	}

	for (int i = 0; i < n_B; i++) {
		ll Sum = B[i];
		for (int j = i + 1; j < n_B; j++) {
			Sum += B[j];
			B.push_back(Sum);
		}
	}
	sort(A.begin(), A.end());
	sort(B.begin(), B.end());

	ll Cnt = 0;
	for (int i = 0; i < A.size(); i++) {
		ll val = A[i];
		if (val >= T) break;
		ll Data = T - val;
		ll y = lower_bound(B.begin(), B.end(), Data) - B.begin();
		ll yy = upper_bound(B.begin(), B.end(), Data) - B.begin();
		Cnt += (yy - y);
	}
	cout << Cnt << endl;
	
	return 0;
}