[C++] 백준 2143번: 두 배열의 합

https://www.acmicpc.net/problem/2143

2143번: 두 배열의 합

Two Pointer로 어떻게 푸는거지 생각하다가 포기하고 답을 찾아봤다. 나는 아무리 생각해도 \(O({N^2M^2})\) 밖에 생각나지 않았다.

https://debuglog.tistory.com/155

백준 2143번 - 두 배열의 합

먼저, 부배열의 합을 미리 전부 구해둔다. 이 때의 Running Time은 \(O(N^2+M^2)\)이다. 그 다음 std::sort 함수를 이용하여 B의 부배열의 합을 정렬한다. 이 때는 \(O(MlgM)\)이 걸린다. 마지막으로 A의 부배열의 합을 처음부터 순회하면서 T와의 차이 만큼의 값이 B의 부배열에 있는지 찾는다. 이 때, upper_bound, lower_bound 함수가 사용된다. 각 함수는 다음과 같은 기능을 한다.

upper_bound: 인수로 넣은 값보다 큰 값이 처음 나오는 위치를 반환함

lower_bound: 인수로 넣은 값과 같거나 그보다 큰 값이 처음 나오는 위치를 반환함

참고: https://chanhuiseok.github.io/posts/algo-55/

알고리즘 - c++ lower_bound, upper_bound 활용하기

따라서, upper_bound로 구한 위치에 lower_bound로 구한 위치를 빼면 해당하는 값의 개수를 구할 수 있다. 그런 값이 없을 경우 반환된 두 값이 같으므로 0이 된다. 알고보면 어렵지 않은 문제지만, 위와 같은 함수를 몰랐었기도 했고 떠올리기엔 쉽지 않았다. 위 방법 말고도 Two Pointer를 이용한 풀이도 있는데, 원리는 비슷하다.

#include <bits/stdc++.h>

using namespace std;

#define TC(T) int T; cin >> T; while(T--)
#define FAST_IO cin.tie(NULL); ios::sync_with_stdio(false);
#define FOR(i, N) for(int i = 0; i < N; ++i)
#define ERR 1e-13
#define EL "\n"
#define VI vector<int>
#define VVI vector<VI>
#define VVVI vector<VVI>
#define VP vector<pair<int, int>>
#define VVP vector<VP>
#define VB vector<bool>


int main() {
#ifdef DEBUG
	freopen("input.txt", "r", stdin);
	//freopen("output.txt", "w", stdout);
#endif
	FAST_IO;

	int T, N, M;
	cin >> T >> N;
	VI A(N, 0);
	FOR(i, N)
		cin >> A[i];
	cin >> M;
	VI B(M, 0);
	FOR(i, M)
		cin >> B[i];

	VI SUMa;
	FOR(i, N) {
		long long sum = 0;
		for (int j = i; j < N; ++j) {
			sum += A[j];
			SUMa.push_back(sum);
		}
	}

	VI SUMb;
	FOR(i, M) {
		long long sum = 0;
		for (int j = i; j < M; ++j) {
			sum += B[j];
			SUMb.push_back(sum);
		}
	}

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

	long long ans = 0;
	for (const auto& sum : SUMa) {
		auto ub = upper_bound(SUMb.begin(), SUMb.end(), T - sum);
		auto lb = lower_bound(SUMb.begin(), SUMb.end(), T - sum);
		ans += ub - lb;
	}
	cout << ans << EL;

	return 0;
}