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

문제 출처

풀이

문제를 보자마자 누적합 문제라는 것은 쉽게 알 수 있다. 하지만 일반적인 누적합 문제로 접근을 하면 시간초과가 발생할 수 밖에 없는데, 그 이유가 n = 1000이라고 했을 때 만들 수 있는 부 배열 쌍의 수가 1000 * 1001 /2가 된다. 즉 대략 50만쯤 되는데 배열이 A,B두 개가 있으므로 둘을 곱하면 어마어마한 시간 복잡도가 발생하게 된다. 따라서 두 개중 하나는 부 배열의 합을 담아두는 map을 사용하였다. 이렇게되면 50만 곱하기 log 50만이라 시간복잡도가 확 낮아지게 된다. 이 외의 별다른 조치는 필요하지 않다. 다음은 전체 코드이다.

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

using namespace std;

long sol(vector<int> &change, vector<int> &nochange, int &T) {
	int n = change.size();
	
	map<int, int> mm;
	
	for ( int i = 0; i < n; i++ ) {
		mm[change[i]]++;
		for ( int j = 0; j < i; j++ ) {
			mm[change[i] - change[j]]++;
		}
	}
	/*for ( auto iter = mm.begin(); iter != mm.end(); iter++ ) {
		cout << "key : " << iter->first << " value : " << iter->second << endl;
	}*/
	int sz = nochange.size();
	long sum = 0;
	int findN = 0;
	for ( int i = 0; i < sz; i++ ) {
		findN = T-nochange[i];
		if ( mm.find(findN) != mm.end() ) {
			sum += mm[findN];
		}
		for ( int j = 0; j < i ; j++ ) {
			findN = T - (nochange[i] - nochange[j]);
			if ( mm.find(findN) != mm.end() ) {
				sum += mm[findN];
			}
		}
	}
	return sum;
}

int main() {
	ios_base::sync_with_stdio(0); cin.tie(NULL); cout.tie(NULL);
	int T;
	cin >> T;
	int n,m;
	cin >> n;
	vector<int> A(n,0);
	for ( int i = 0; i < n; i++ ) {
		int t;
		cin >> t;
		if ( i == 0 ) {
			A[i] = t;
			continue;
		}
		A[i] = A[i - 1] + t;
	}
	cin >> m;
	vector<int> B(m,0);
	for ( int i = 0; i < m; i++ ) {
		int t;
		cin >> t;
		if ( i == 0 ) {
			B[i] = t;
			continue;
		}
		B[i] = B[i - 1] + t;
	}
	long res = 0;
	if ( n < m ) {
		res = sol(B, A, T);

	}
	else {
		res = sol(A, B, T);
	}
	cout << res << endl;
	return 0;
}