[백준 알고리즘][C++] 이분 탐색(binary search) - 2143번 두 배열의 합

안녕하세요~

​

이번시간에는 2143번 두 배열의 합 문제를 풀어보려고 합니다.

 

www.acmicpc.net/problem/2143

2143번: 두 배열의 합

 

명심하셔야 될 것은 일단 문제를 풀기 전에 코딩을 먼저 시작하지 마시고, 반드시 연습장으로 생각을 정리하고! 이 후에 생각이 맞는 것 같다고 확신이 되면 그 때 코딩을 시작하세요. 코드를 작성하면서 문제를 풀게되면 정리가 안된 상태기 때문에 코드 수정이 잦아지고, 정신이 없어지게 됩니다ㅠㅠ

그러니 반드시 노트에 어떻게 풀 것인지에 대한 정리를 다 하시고 코딩을 시작하시기 바랍니다!!

​

binary search 공부에 좋은 문제입니다.

 

문제

한 배열 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

정답 코드

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;
int A[1001];
int B[1001];
vector<long long>sum_A;
vector<long long>sum_B;
int N, M, T;
int main()
{
	cin >> T;
	cin >> N;
	for (int i = 0; i < N; i++)
	{
		cin >> A[i];
	}
	cin >> M;
	for (int i = 0; i < M; i++)
	{
		cin >> B[i];
	}
	long long sum = 0;
	//i번째 부터 만들 수 있는 합을 모두 더함
	for (int i = 0; i < N; i++)
	{
		sum = A[i];
		sum_A.push_back(sum);
		for (int j = i + 1; j < N; j++)
		{
			sum = sum + A[j];
			sum_A.push_back(sum);
		}
	}
	for (int i = 0; i < M; i++)
	{
		sum = B[i];
		sum_B.push_back(sum);
		for (int j = i + 1; j < M; j++)
		{
			sum = sum + B[j];
			sum_B.push_back(sum);
		}
	}
	sort(sum_A.begin(), sum_A.end());
	sort(sum_B.begin(), sum_B.end());
	long long count = 0;
	for (int i = 0; i < sum_A.size(); i++)
	{
		int start = 0;
		int end = sum_B.size() - 1;
		int target = sum_A[i];
		while (start <= end)
		{
			int mid = (start + end) / 2;
			if ((T - target) >= sum_B[mid])
			{
				start = mid + 1;
			}
			else
			{
				end = mid - 1;
			}
		}
		int upper_bound = start;
		start = 0;
		end = sum_B.size() - 1;
		while (start <= end)
		{
			int mid = (start + end) / 2;
			if ((T - target) > sum_B[mid])
			{
				start = mid + 1;
			}
			else
			{
				end = mid - 1;
			}
		}
		int lower_bound = start;

		count = count + (upper_bound - lower_bound);
	}
	cout << count << endl;
}