백준2143 두 배열의 합

작성일 | In

문제 링크

  • http://icpc.me/2143

문제 출처

  • 2001 전국 본선 고등부1

사용 알고리즘

  • 이진 탐색

시간복잡도

  • O(N2logN2)

풀이

A 배열의 모든 부 배열의 합과, B배열의 모든 부 배열의 합을 먼저 구해줍시다.
A와 B의 prefix sum을 구해놓았다면, O(N2)만에 각각의 부배열의 합을 모두 구할 수 있습니다. 그 배열을 각각 aa, bb라고 합시다.

만약 aa의 어떤 원소 i가 bb의 어떤 원소 j의 합이 t가 되기 위해서는, i+j=t를 만족해야 합니다. 당연하게도, j=t-i를 만족하겠죠.
aa의 원소 i가 있을 때, bb에 있는 t-i값의 개수의 합이 정답이 되겠네요.

aa와 bb를 정렬해준 다음에, aa의 모든 원소 i에 대해 bb에서 t-i의 개수를 구해줍시다.

전체 코드

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
#include <bits/stdc++.h>
using namespace std;

typedef long long ll;

ll t;
int n, m;
ll a[1010], b[1010];
ll sa[1010], sb[1010];

vector<ll> aa, bb;

int main(){
	ios_base::sync_with_stdio(0); cin.tie(0);
	cin >> t;
	cin >> n;
	for(int i=1; i<=n; i++){
		cin >> a[i]; sa[i] = sa[i-1] + a[i];
	}
	cin >> m;
	for(int i=1; i<=m; i++){
		cin >> b[i]; sb[i] = sb[i-1] + b[i];
	}

	for(int i=1; i<=n; i++){
		for(int j=i; j<=n; j++){
			aa.push_back(sa[j] - sa[i-1]);
		}
	}
	for(int i=1; i<=m; i++){
		for(int j=i; j<=m; j++){
			bb.push_back(sb[j] - sb[i-1]);
		}
	}

	ll ans = 0;
	sort(aa.begin(), aa.end());
	sort(bb.begin(), bb.end());
	for(int i=0; i<aa.size(); i++){
		auto it = equal_range(bb.begin(), bb.end(), t - aa[i]);
		ans += distance(it.first, it.second);
	}
	cout << ans;
}