[BOJ/Gold 3] 백준 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을 출력한다.

예제 입력 1

5
4
1 3 1 2
3
1 3 2

예제 출력 1

7

알고리즘 분류

• 이분 탐색
• 누적 합

풀이

배열 A, B의 누적 합이 될 수 있는 모든 수를 구해주고 각각 벡터의 원소로 push한다.

그리고 배열 A의 누적 합의 벡터 SumA의 각각의 원소마다, T에서 원소를 뺀 값이 될 수 있는 SumB의 원소들을 이분 탐색을 통해 찾는다. (lower_bound : 어떤 값과 같은 원소의 인덱스 중 가장 작은 값, upper_bound : 어떤 값보다 큰 원소의 인덱스 중 가장 작은 값)

코드

#include <iostream>
#include <string>
#include <cmath>
#include <unordered_set>
#include <unordered_map>
#include <bitset>
#include <ctime>
#include <stack>
#include <queue>
#include <vector>
#include <algorithm>
#include <climits>
#define FASTIO cin.tie(NULL); cout.tie(NULL); ios::sync_with_stdio(false);
#define MAX 1001
#define LL long long
#define INF 1e9

using namespace std;
int T, N, M;
int A[MAX], B[MAX];
vector<int> SumA, SumB;
LL Answer = 0;

void Input() {
	cin >> T;
	cin >> N;
	for (int i = 1; i <= N; i++) {
		cin >> A[i];
	}
	cin >> M;
	for (int i = 1; i <= M; i++) {
		cin >> B[i];
	}
}

void Settings() {
	for (int i = 1; i <= N; i++) {
		int Tmp = A[i];
		SumA.push_back(Tmp);
		for (int j = i + 1; j <= N; j++) {
			Tmp += A[j];
			SumA.push_back(Tmp);
		}
	}
	for (int i = 1; i <= M; i++) {
		int Tmp = B[i];
		SumB.push_back(Tmp);
		for (int j = i + 1; j <= M; j++) {
			Tmp += B[j];
			SumB.push_back(Tmp);
		}
	}
	sort(SumB.begin(), SumB.end());
	for (int i = 0; i < SumA.size(); i++) {
		int Rest = T - SumA[i];
		int Up = upper_bound(SumB.begin(), SumB.end(), Rest) - SumB.begin();
		int Same = lower_bound(SumB.begin(), SumB.end(), Rest) - SumB.begin();
		Answer += Up - Same;
	}
}

void Find_Answer() {
	cout << Answer << "\n";
}

int main() {
	FASTIO
	Input();
	Settings();
	Find_Answer();

	return 0;
}