[Baekjoon] 2143번: 두 배열의 합

🥅 문제

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

 

🔍 설계

부배열의 형태를 보면, 누적 합을 이용해야 한다는 것은 쉽게 파악할 수 있다. 

 

처음에는 A와 B의 부배열을 구해서 서로 다른 맵에 저장한 후, 한 쪽에서 반복자를 사용해 반복하면서, (T - 한 쪽에 있는 부배열의 합)이 다른 쪽에 있으면, (한 쪽에 있는 부배열의 합의 개수) * (다른 쪽에 있는 부배열의 합의 개수)로 구하려고 했으나 메모리 초과가 발생하였다. 맵 하나가 최대 1000000개 * 2 * 4Bytes 이므로 64MB를 초과하기 때문이었다.

 

그래서, 맵을 하나만 만들고, 다른 쪽에서 부배열의 합을 만들면서 바로 찾는 방식을 사용하였다. 부배열의 합을 만들기 위해 이중 반복문을 사용하기 때문에 1000000, 한 쪽에서 T - (만든 부배열의 합)을 찾기 위해 log1000 이므로 10^7이기 때문에 시간 초과가 발생하지 않는다.

---

 

 

위와 마찬가지로, 한 쪽에서 부배열의 합을 만들고, 다른 쪽에 T - 부배열의 합이 몇 개가 있는지 구해서 결과에 더하면 된다. 그 개수는 lower_bound와 upper_bound를 이용하여 위 방법보다 빠르게, 메모리를 적게 사용하여 구할 수 있다.

 

upper_bound는 key보다 큰 최초의 위치를 반환하고, lower_bound는 key보다 크거나 같은 최초의 위치를 반환한다. upper_bound - lower_bound는 배열에서 key의 개수가 된다는 것을 이용한다.

 

👊 풀이

#include <iostream>
#include <vector>
#include <map>
#include <algorithm>
using namespace std;

int t, n, m;
vector<int> a(1000);
vector<int> aSum(1000);
vector<int> b(1000);
vector<int> bSum(1000);
map<int, int> aMap;
long long result = 0;

int main(void)
{
	ios_base::sync_with_stdio(false);
	cin.tie(nullptr);

	cin >> t;
	cin >> n;
	for (int i = 0; i < n; i++) {
		cin >> a[i];
		if (i > 0) {
			aSum[i] = a[i] + aSum[i - 1];
		}
		else {
			aSum[i] = a[i];
		}
	}
	cin >> m;
	for (int i = 0; i < m; i++) {
		cin >> b[i];
		if (i > 0) {
			bSum[i] = b[i] + bSum[i - 1];
		}
		else {
			bSum[i] = b[i];
		}
	}

	for (int i = 0; i < n; i++) {
		if (aMap.find(aSum[i]) != aMap.end()) {
			aMap[aSum[i]] += 1;
		}
		else {
			aMap.insert({ aSum[i] , 1 });
		}

		for (int j = 0; j < i; j++) {
			long long part_sum = aSum[i] - aSum[j];
			if (aMap.find(aSum[i] - aSum[j]) != aMap.end()) {
				aMap[part_sum] += 1;
			}
			else {
				aMap.insert({ part_sum, 1 });
			}
		}
	}

	for (int i = 0; i < m; i++) {
		if (aMap.find(t - bSum[i]) != aMap.end()) {
			result += aMap[t - bSum[i]];
		}
		for (int j = 0; j < i; j++) {
			long long part_sum = bSum[i] - bSum[j];
			if (aMap.find(t - part_sum) != aMap.end()) {
				result += aMap[t - part_sum];
			}
		}
	}

	cout << result << "\n";

	return 0;
}

 

 

👊 풀이2

#include <iostream>
#include <vector>
#include <map>
#include <algorithm>
using namespace std;

int t, n, m, part_sum;
vector<int> a(1000);
vector<int> aSum(1000);
vector<int> b(1000);
vector<int> bSum(1000);
vector<int> partSumA;
long long result = 0;

int main(void)
{
	ios_base::sync_with_stdio(false);
	cin.tie(nullptr);

	cin >> t;
	cin >> n;
	for (int i = 0; i < n; i++) {
		cin >> a[i];
		if (i > 0) {
			aSum[i] = a[i] + aSum[i - 1];
		}
		else {
			aSum[i] = a[i];
		}
	}
	cin >> m;
	for (int i = 0; i < m; i++) {
		cin >> b[i];
		if (i > 0) {
			bSum[i] = b[i] + bSum[i - 1];
		}
		else {
			bSum[i] = b[i];
		}
	}

	for (int i = 0; i < n; i++) {
		partSumA.push_back(aSum[i]);
		for (int j = 0; j < i; j++) {
			partSumA.push_back(aSum[i] - aSum[j]);
		}
	}

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

	for (int i = 0; i < m; i++) {
		result += upper_bound(partSumA.begin(), partSumA.end(), t - bSum[i]) -lower_bound(partSumA.begin(), partSumA.end(), t - bSum[i]);
		for (int j = 0; j < i; j++) {
			part_sum = bSum[i] - bSum[j];
			result += upper_bound(partSumA.begin(), partSumA.end(), t - part_sum) - lower_bound(partSumA.begin(), partSumA.end(), t - part_sum);
		}
	}

	cout << result << "\n";

	return 0;
}

 

📝 메모

첫 인덱스부터 어느 인덱스까지의 부분합을 넣는 과정에서 그 값이 이미 존재할 수 있다.