[정올/JUNGOL] 1357번 합이 0이 되는 4개의 숫자들 (C/C++)

 정보 올림피아드 알고리즘 1357번 합이 0이 되는 4개의 숫자들
https://www.jungol.co.kr/problem/1357

JUNGOL

* 사용언어 : C언어, C++

 

1. 문제

최대 4,000개 정수로 이루어진 4개의 배열이 주어짐
각 배열의 값을 1개 씩 사용했을 때 네 수의 합이 0이 되는 쌍의 개수를 출력

 

2. 풀이

최대 4,000개의 정수로 이루어져있기 때문에 for 문 4개로 전체 탐색하면
4,000 ^ 4 = 256,000,000,000,000 회 반복해야 합니다.
rough 하게 초당 1억 번, 아니 10억 번 반복 가능하다고 해도 제한 시간 안에 절대 풀 수 없습니다.

방법을 못 찾아서 고민하다가 구글링해서 풀었습니다.

1. A 와 B 로 가능한 모든 합을 배열(ab) 로 구성 (4,000 * 4,000)
2. C 와 D 로 가능한 모든 합을 배열(cd) 로 구성 (4,000 * 4,000)
3. ab 와 cd 배열을 정렬 (저는 개폐구간 merge sort 로 정렬했습니다)
4. ab 는 0 부터 15,999,999 까지, cd 는 15,999,999 부터 0 까지 탐색하면서 합이 0이 되는 쌍을 찾습니다.
5. 한 배열에 같은 숫자가 여러 개 있는 경우 해당 숫자의 개수를 곱해야 합니다.
   (ex. ab = {-7, -7, -7, -1}, cd = {0, 7, 7, 8} 인 경우 ab 의 -7이 3개, cd 의 7이 2개이므로 3 * 2 = 6개)

나머지는 풀이는 아래 코드 참고하시면 됩니다.

 

3. 코드

#ifndef _CRT_SECURE_NO_WARNINGS
#define _CRT_SECURE_NO_WARNINGS
#endif

#include <cstdio>

using LL = long long;
const int LM = 4000;
int arr[4][LM];
int ab[LM * LM], cd[LM * LM], tmp[LM * LM];

void mergeSort(int* A, int s, int e) { // [s, e)
	if (s + 1 >= e) return;

	int m = (s + e) / 2;
	mergeSort(A, s, m);
	mergeSort(A, m, e);

	int i = s, j = m, k = s;
	while (i < m && j < e) {
		if (A[i] <= A[j]) tmp[k++] = A[i++];
		else tmp[k++] = A[j++];
	}
	while (i < m) tmp[k++] = A[i++];
	while (j < e) tmp[k++] = A[j++];

	for (i = s; i < e; ++i) A[i] = tmp[i];
}

int main() {
#ifdef _WIN32
	freopen("1357.txt", "r", stdin);
#endif // _WIN32
	int n, i, j;
	scanf("%d", &n);
	for (i = 0; i < n; ++i) scanf("%d %d %d %d", &arr[0][i], &arr[1][i], &arr[2][i], &arr[3][i]);

	int idx = 0;
	for (i = 0; i < n; ++i) {
		for (j = 0; j < n; ++j) {
			ab[idx] = arr[0][i] + arr[1][j];
			cd[idx] = arr[2][i] + arr[3][j];
			++idx;
		}
	}

	mergeSort(ab, 0, n * n);
	mergeSort(cd, 0, n * n);

	i = 0, j = n * n - 1;
	LL cnt_ab, cnt_cd, res = 0;
	while (i < n * n && j >= 0) {
		if (ab[i] + cd[j] < 0) ++i;
		else if (ab[i] + cd[j] > 0) --j;
		else {
			cnt_ab = 1, cnt_cd = 1;
			while (ab[i] == ab[i + cnt_ab] && i + cnt_ab < n * n) ++cnt_ab;
			while (cd[j] == cd[j - cnt_cd] && j - cnt_cd >= 0) ++cnt_cd;

			res += cnt_ab * cnt_cd;
			i += cnt_ab;
			j -= cnt_cd;
		}
	}

	printf("%lld\n", res);
	return 0;
}

* case 에 따라서 res 가 int 범위를 넘어갈 수 있으니 long long 으로 저장 및 출력합니다. (ex. 4,000개 모두 0 0 0 0인 케이스)