[백준] 두 배열의 합 (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을 출력한다.

---

합이 0인 네 정수(7453번)와 비슷한 문제이다.

두 배열의 연속적인 부분 배열의 합을 통해 T를 만들 수 있는 경우의 수를 구하는 법이다.

A와 B 배열에 대해 가능한 연속합을 모두 구해서 aSum, bSum 배열에 넣어놓자.

그 후 T - aSum[]와 같은 값을 가지는 bSum의 개수를 구하면 된다.

마찬가지로 lower, upper bound를 이용해서 문제를 해결하자.

 

코드 원본 : https://github.com/chosh95/STUDY/blob/master/BaekJoon/2020/%EB%91%90%20%EB%B0%B0%EC%97%B4%EC%9D%98%20%ED%95%A9%20(2143%EB%B2%88).cpp

chosh95/STUDY

C++ 코드

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;
int T, N, M;
int A[1001], B[1001];

int main()
{
	cin >> T >> N;
	for (int i = 0; i < N; i++) cin >> A[i];
	cin >> M;
	for (int i = 0; i < M; i++) cin >> B[i];

	vector<int> aSum, bSum;
	
	for (int i = 0; i < N; i++) {
		int sum = A[i];
		aSum.push_back(sum);
		for (int j = i + 1; j < N; j++) {
			sum += A[j];
			aSum.push_back(sum);
		}
	}

	for (int i = 0; i < M; i++) {
		int sum = B[i];
		bSum.push_back(sum);
		for (int j = i + 1; j < M; j++) {
			sum += B[j];
			bSum.push_back(sum);
		}
	}

	sort(bSum.begin(), bSum.end());
	long long ans = 0;
	for (int i = 0; i < aSum.size(); i++) {
		int target = T - aSum[i];
		int lo = lower_bound(bSum.begin(), bSum.end(), target) - bSum.begin();
		int hi = upper_bound(bSum.begin(), bSum.end(), target) - bSum.begin();
		ans += (hi - lo);
	}

	cout << ans;
}