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

 

 

누적 합과 이분 탐색으로 쉽게 풀 수 있는 문제

lower_bound, upper_bound 오랜만에 직접 구현해보다가 결국 귀찮아서 라이브러리에 있는 거 썼다..

 

풀이 과정

1. 누적 합으로 A, B 배열의 부 배열의 합이 될 수 있는 경우의 수를 모두 vector에 저장해준다. (최대 500500 경우의 수)
2. 둘 중 하나를 sort 해준다.
3. upper_bound - lower_bound로 target(t - vec_a[i])이 몇 개나 있는 지 구해서 dap에 더해준다.
   • dap은 int 범위를 넘을 수 있으므로 long이나 long long으로 처리해준다.
   • 못 찾았을 경우 up, lo 모두 end() 저장소의 위치를 가르키므로 예외처리 안 해줘도 0을 더한다.

 

#include<iostream>
#include<algorithm>
#include<vector>
using namespace std;
typedef long long ll;

int a[10001], b[10001];
vector<int> vec_a, vec_b;

int main()
{
	int n, m, t;
	scanf("%d", &t);

	scanf("%d", &n);
	for (int i = 0; i < n; i++)scanf("%d", &a[i]);
	scanf("%d", &m);
	for (int i = 0; i < m; i++)scanf("%d", &b[i]);

	for (int i = 0; i < n; i++) {
		int sum = 0;
		for (int j = i; j < n; j++) {
			sum += a[j];
			vec_a.push_back(sum);
		}
	}

	for (int i = 0; i < m; i++) {
		int sum = 0;
		for (int j = i; j < m; j++) {
			sum += b[j];
			vec_b.push_back(sum);
		}
	}

	int a_size = vec_a.size();
	sort(vec_b.begin(), vec_b.end());
	
	ll dap = 0;
	for (int i = 0; i < a_size; i++) {
		int up = upper_bound(vec_b.begin(), vec_b.end(), t - vec_a[i]) - vec_b.begin();
		int lo = lower_bound(vec_b.begin(), vec_b.end(), t - vec_a[i]) - vec_b.begin();
		dap += up - lo;
	}
	printf("%lld\n", dap);
	return 0;
}