[Baekjoon, 2143] 두 배열의 합

두 배열에 대한 모든 부분합을 미리 구한 다음, 한 배열을 정렬하고 이분탐색을 이용하여 답의 갯수를 찾으면 되는 문제이다.

나는 오랜만에 접한 lower_bound(), upper_bound() 함수가 정확히 어떤값을 반환하는지 헷갈렸다. 그래서 정리하고자 이 글을 쓰게되었다.

• lower_bound(): 찾으려는 값중 제일 첫 인덱스의 이터레이터를 반환한다. 없다면 찾으려는 값보다 큰 값중 가장 작은 값의 이터레이터를 반환한다.
• upper_bound(): 찾으려는 값보다 큰 값중 가장 작은 값의 이터레이터를 반환한다.

위의 함수를 이용하여 미리 구한 부분합에서 T = subsum1[i] + subsum2[j]를 만족하게 하는 구간의 크기를 모두 더하면 된다.

구현

#include <iostream>
#include <vector>
#include <queue>
#include <map>
#include <algorithm>
#include <cstring>
#include <stack>
#include <cmath>
#include <climits>
#include <cctype>
using namespace std;

const int MAX = 1000 + 1;

int T, N, M;

vector<long long> calcSubsum(int arr[], int len) {
	vector<long long> ret;
	for (int i = 0; i < len; ++i) {
		long long sum = arr[i];
		ret.push_back(sum);
		for (int j = i + 1; j < len; ++j) {
			sum += arr[j];
			ret.push_back(sum);
		}
	}
	return ret;
}

int main() {
	ios_base::sync_with_stdio(0); cin.tie(0);

	int A[MAX], B[MAX];

	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];
	}

	vector<long long> subsum1 = calcSubsum(A, N);
	vector<long long> subsum2 = calcSubsum(B, M);

	sort(subsum2.begin(), subsum2.end());

	long long ret = 0;
	for (int i = 0; i < subsum1.size(); ++i) {
		int lo = lower_bound(subsum2.begin(), subsum2.end(), T - subsum1[i]) - subsum2.begin();
		int hi = upper_bound(subsum2.begin(), subsum2.end(), T - subsum1[i]) - subsum2.begin();

		ret += hi - lo;
	}

	cout << ret << endl;

	return 0;
}