백준 7453 : 합이 0인 네 정수

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

7453번: 합이 0인 네 정수

어떤 배열 A,B,C,D가 주어질 때 각각의 배열에서 하나의 원소를 골라 네 수를 더했을 때 합이 0이 되는 경우의 수를 구하는 문제이다.

 

각각의 배열 사이즈는 최대 4000이고, 문제에서 주어진 시간은 12초이므로 O(N^2) 혹은 O(N^2 logN)의 시간 복잡도로 문제를 해결해야 한다. 나는 두 가지 방법으로 문제를 해결했는데 기본적인 개념은 다음과 같다.

 

일단 배열 네 개의 원소를 하나씩 골라 합을 구하는 시간 복잡도는 O(N^4)이므로 배열을 두개 두개씩 묶어서 prefix sum 연산을 수행해서 두 개의 배열로 줄였다. 물론 여기서 다시 완전탐색을 돌리면 O(N^4)로 같은 결과가 나오지만 이분탐색을 이용하면 배열을 합치는데 O(N^2), 탐색에 O(N^2 log N^2)이 되므로 시간 내에 문제 해결이 가능하다.

해쉬맵 자료구조를 사용 시 더 빠르게 문제 해결이 가능하지만 이 문제에서는 해쉬맵 자료구조 사용 시 메모리 초과가 발생해 사용할 수 없다.

 

또 하나의 방법은 배열을 두 개씩 묶는 방법까지는 동일하고, 이후 투포인터로 푸는 방법이다. 이 방법의 경우 배열을 합치는데 O(N^2), 탐색에 O(N^2)이 소요되므로 이분탐색 방법보다 더 효율적인 풀이라고 할 수 있다.

 

먼저 이분탐색 코드다.

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

using namespace std;

vector<int> prefix_cd;

int n;
int arr[4][4000];

void prefixsum() {
	int x = 2, y = 3;
	for (int i = 0; i < n; i++) {
		for (int j = 0; j < n; j++) {
			prefix_cd.emplace_back(arr[x][i] + arr[y][j]);
		}
	}
}

long long solution() {
	long long cnt = 0;

	prefixsum();

	//이분 탐색을 위한 정렬
	sort(prefix_cd.begin(), prefix_cd.end());

	int target;
	for (int i = 0; i < n; i++) {
		for (int j = 0; j < n; j++) {
			target = arr[0][i] + arr[1][j];
            //이분탐색을 사용하는 upper_bound, lower_bound 내장 함수 사용
			cnt += (upper_bound(prefix_cd.begin(), prefix_cd.end(), -target) - lower_bound(prefix_cd.begin(), prefix_cd.end(), -target));
		}
	}

	return cnt;
}

int main() {
	ios::sync_with_stdio(0), cin.tie(0), cout.tie(0);
	cin >> n;
	for (int i = 0; i < n; i++) {
		for (int j = 0; j < 4; j++) {
			cin >> arr[j][i];
		}
	}

	cout << solution();

	return 0;
}

메모리: 100404 kb

시간: 11488 ms

 

다음은 투포인터 풀이다.

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

using namespace std;

vector<int> prefix_ab;
vector<int> prefix_cd;

int n;
int arr[4][4000];

void prefixsum() {
	int x = 0, y = 1;
	for (int i = 0; i < n; i++) {
		for (int j = 0; j < n; j++) {
			prefix_ab.emplace_back(arr[x][i] + arr[y][j]);
		}
	}
		
	x = 2, y = 3;
	for (int i = 0; i < n; i++) {
		for (int j = 0; j < n; j++) {
			prefix_cd.emplace_back(arr[x][i] + arr[y][j]);
		}
	}
}

long long solution() {
	long long cnt = 0;
	long long same_lv, same_rv;
	int lidx, ridx, lv, rv;

	prefixsum();

	//투포인터 사용을 위한 정렬
	sort(prefix_ab.begin(), prefix_ab.end());
	sort(prefix_cd.begin(), prefix_cd.end());

	lidx = 0, ridx = prefix_cd.size() - 1;
	while (lidx < prefix_ab.size() && ridx >= 0) {
		rv = prefix_cd[ridx];
		lv = prefix_ab[lidx];

		if (rv + lv == 0) {
			same_lv = 1, same_rv = 1;
            //합이 0인 경우 두 배열에서 같은 값들의 개수의 곱을 cnt에 더해줌
			while (++lidx < prefix_ab.size() && prefix_ab[lidx] == lv) ++same_lv;
			while (--ridx >= 0 && prefix_cd[ridx] == rv) ++same_rv;
			cnt += same_lv * same_rv;
		}
		else if (rv + lv > 0) --ridx;
		else ++lidx;
	}
    
	return cnt;
}

int main() {
	ios::sync_with_stdio(0), cin.tie(0), cout.tie(0);
	cin >> n;
	for (int i = 0; i < n; i++) {
		for (int j = 0; j < 4; j++) {
			cin >> arr[j][i];
		}
	}

	cout << solution();

	return 0;
}

메모리: 166060 kb

시간: 3056 ms

 

시간을 보면 투포인터 풀이가 훨씬 효율적인 것을 알 수 있다.