백준2143 두 배열의 합 - C++

두 배열을 주고, 두 배열의 연속된 부분집합의 합이 

특정 숫자 T를 만족하는 개수를 찾으란다..

 

T 는 -10억 ~ 10억이고 

A, B 배열은 1000개 까지 있다. 각 배열의 값은 -1백만 ~ 1백만이다.

 

일단 배열의 부분집합수만 해도 Sigma(1000) 이니 각 50만개가 넘는다.

두 개 결합하는 경우의 수는 50만 x 50만이므로 2초 안에 불가능하므로 

필요한 경우만 빠르게 찾는 것이 목표다.

 

나는 두 가지를 이용해서 간편하게 풀길 원했다.

 

1. presum

2. lower bound, upper bound

 

문제에서 연속된 부분집합만을 요구하므로

미리 구간합을 계산하여 가지고 있을 수 있다.

 

예를 들어, A 가 [1,2,3,4,5] 라면

presum = [1, 2, 3, 4, 5, 1+2, 2+3, 3+4, 4+5,1+2+3,2+3+4,3+4+5,1+2+3+4,2+3+4+5,1+2+3+4+5] 

라는 부분수열의 합을 미리 구할 수 있다.

이 상태에서 부분수열이 어디부터 어디까지인지 알 필요가 없고

조합만 계산하면 된다.

 

예를 들어 B 도 [1,2,3,4,5] 라면

presum A, B = [1,2,3,4,5,3,5,6,9,6,9,12,10,14,15] 를 구할 수 있고

이를 정렬하여 [1,2,3,3,4,5,5,6,6,9,9,10,12,14,15] 를 얻을 수 있다.

 

presum A + presum B = T 가 되는 수를 찾을 수 있다. 

가령 T 가 13이라면

presum A = 1, presum B = 12

presum A = 3, presum B = 10

presum A = 4, presum B = 9

presum A = 9, presum B = 4

presum A = 10, presum B = 3

presum A = 12, presum B = 1

을 구할 수 있다.

 

이 때 presum 에 저장된 값들은 모두 (유니크한 부분수열)의 합이기 때문에

각 조합은 모두 부분수열의 합의 조합이므로 

성공할 때마다 개수를 세어주면 되겠다.

 

근데 또 하나씩 세어주기 싫다. 몇 개가 될 지 모르기 때문에.

그래서 upper bound 와 lower bound 를 구해서 개수를 구해준다.

가령 presum A = 3, presum B = 10 인 경우 개수는

n(presum A = 3) = 2 개, n(presum B = 10) 1개 이므로 2개의 경우가 있다.

즉 upper_bound - lower_bound 의 곱으로 이루어져 있다.

이것을 이용해서 모든 경우를 탐색하는 것을 피할 수 있다.

(근데 곱하는게 잘 안 먹어서 v2 에만 적용했다.)

 

아래의 코드는 N^2 개의 개수를 2번 quick sort(std::algorithm 의 sort 는 quick sort임이 알려져있다.)

하고, N^2 개를 순회하면서 upper_bound 와 lower_bound 를 각각 구하므로

시간복잡도는 2 * N^2logN^2 + 2 * N^2 * logN^2 = O(N^2 log N) 이 되겠다.  

#include <iostream>
#include <vector>
#include <algorithm>

using namespace std;
typedef long long ll;
int A[1001];
int B[1001];
vector<ll> v1;
vector<ll> v2;



int main()
{
	cin.tie(0);
	ios_base::sync_with_stdio(0);
	ll T;
	cin >> T;
	int n;
	cin >> n;
	for (int i = 0; i < n; i++)
	{
		int a;
		cin >> a;
		A[i] = a;
		v1.push_back(a);
	}
	int m;
	cin >> m;
	for (int i = 0; i < m; i++)
	{
		int a;
		cin >> a;
		B[i] = a;
		v2.push_back(a);
	}

	for (int i = 0; i < n; i++)
	{
		ll sum;
		sum = A[i];
		for (int j = i + 1; j < n; j++)
		{
			sum += A[j];
			v1.push_back(sum);
		}

	}
	for (int i = 0; i < m; i++)
	{
		ll sum;
		sum = B[i];
		for (int j = i + 1; j < m; j++)
		{
			sum += B[j];
			v2.push_back(sum);
		}
	}
	sort(v1.begin(), v1.end());
	sort(v2.begin(), v2.end());
	long long result = 0;
	for (int i = 0; i < v1.size(); i++)
	{
		int low = lower_bound(v2.begin(), v2.end(), T - v1[i])-v2.begin();
		int high = upper_bound(v2.begin(), v2.end(), T - v1[i])-v2.begin();

		result += high - low;
	}
	cout << result;
}