[Java] BOJ7453, 백준 - 합이 0인 네 정수

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

7453번: 합이 0인 네 정수

 

시간 제한	메모리 제한	제출	정답	맞은 사람	정답 비율
12 초	1024MB	22199	5386	3031	22.519%

문제

정수로 이루어진 크기가 같은 배열 A, B, C, D가 있다.

A[a], B[b], C[c], D[d]의 합이 0인 (a, b, c, d) 쌍의 개수를 구하는 프로그램을 작성하시오.

 

입력

첫째 줄에 배열의 크기 n (1 ≤ n ≤ 4000)이 주어진다. 다음 n개 줄에는 A, B, C, D에 포함되는 정수가 공백으로 구분되어져서 주어진다. 배열에 들어있는 정수의 절댓값은 최대 228이다.

 

출력

합이 0이 되는 쌍의 개수를 출력한다.

 

<풀이>

처음엔 그냥 브루트포스로 풀려고 했지만, N의 범위가 4000이기 때문에 N^4하면 범위가 어마무시해져버린다.

그래서 ab와 cd로 나눈다음에, 값을 비교하기로 한다.

ab의 합 구하기 : n^2

cd의 합 구하기 :  n^2

 

나누고 난 이후에 전부 비교한다고 가정하면 역시 시간복잡도가 n^4이기 때문에, 조금 더 효율적인 비교 방법을 고려해보도록 하겠다.

우리는 ab와 cd를 합쳐서 0이 되는 숫자를 찾으려고 한다.

그런데 ab가 음수이고, cd역시 음수라면 당연히 합쳐서 0이 나올 수 없을 것이다.

이 아이디어에서 시작해보았다.

우선 음수/0/양수를 가장 효율적으로 비교할 수 있는 방법? 바로 정렬이다.

정렬을 하면 순차적으로 숫자를 비교하면 되기 때문에 위를 가장 효율적으로 비교할 수 있을 것이다.

자바에서 제공하는 Arrays.sort(더블 피봇 퀵 정렬)을 이용하면 시간복잡도는 O(n log(n))

두 번 실행하기 때문에

2*4000^2*log(4000^2)

약 2억3천

 

예제를 이용해서 예시를 들어보면 정렬 이후에는 다음과 같은 꼴이 될 것이다.

 

<초기상태>

a	b	c	d
-45	22	42	-16
-41	-27	56	30
-36	53	-37	77

<ab합과 cd합으로 나눈 상태>

ab	cd
-23	26
-72	72
8	119
-19	40
-68	86
12	133
-14	-53
-63	-7
17	40

<정렬>

ab	cd
-72	-53
-68	-7
-63	26
-23	40
-19	40
-14	72
8	86
12	119
17	133

 

이제 ab와 cd의 합을 비교해주면 되는데, 위에서 설명했던 방법처럼 ab가 음수이면 cd가 양수인 경우를 계산해야 하기 때문에, ab는 0번째 인덱스부터/cd는 맨 마지막 인덱스부터 합쳐주도록 하겠다.

그러면 연산 순서를 살펴보자

1. ab=-72일 때
   1. -72+133 > 0   넘어간다.
   2. -72+119 > 0   넘어간다.
   3. -72 + 86 > 0   넘어간다.
   4. -72 + 72 = 0   카운트
   5. -72 + 40 < 0   넘어간다.
   6. ... -72 + 26 < 0 넘어간다.
   7. cd가 -7로, 음수가 되어 종료.
2. ab=-68일 때
   1. -68+133 > 0   넘어간다.
   2. -68+119 > 0   넘어간다.
   3. 잠깐, 뭔가 이상하지 않은가?
   4. ab도 정렬을 했기 때문에 ab의 인덱스가 진행하면 ab 값도 커지기 때문에 이전에 -72에서 넘어왔다면 cd에서도 72보다 작은 부분부터만 검사를 해주면 되는데, 지금 쓸데없이 반복적인 연산이 반복되고 있다.

따라서 ab와 cd를 각각 음수/양수로 나누어 생각하기보다는

ab+cd가 음수인지, 양수인지 나누어서 생각해보도록 하자(==0일때는 정답으로 카운트)

그러면 다음과 같은 구조가 완성된다.

// sum은 ab와 cd의 합
if (sum == 0) { // 정답을 찾았을 경우
	// 카운트한다
} else if (sum > 0) { 
	// 0보다 클 경우 cd의 인덱스값을 줄인다
} else { 
	// 0보다 작을 경우 ab의 인덱스값을 높인다
}

그렇다면 정답을 찾았을 경우 카운트한다음 인덱스는 어떻게 진행시켜야 할까?

만약 그냥 둘다 하나씩 진행시켜버린다면 다음과 같은 상황에선 문제가 발생한다.

ab	cd
-11	76
-3	3
2	3
99	-9

실제로는 ab에서의 -3이 cd에서의 3과 더해져서 0이 되는 경우가 두 번이지만,

인덱스를 한꺼번에 증가시켜버린다면 한 번만 카운트될것이다.

이렇게 ab든 cd든 같은 수가 중복해서 나오는 경우에 문제가 발생하게된다.

그래서 합이 0이 되는 수를 발견했을 때는, 중복된 값이 없을때까지 계속해서 찾아줘야한다.

// 중복된 값이 없을때까지 탐색
while (l_cursor < n * n && ab[l_cursor] == abVal) {
	l_cursor++;
	abcnt++;
}
while (r_cursor >= 0 && cd[r_cursor] == cdVal) {
	r_cursor--;
	cdcnt++;
}
cnt += (abcnt * cdcnt);

중복된 값이 있었다면 서로 곱해서 카운트해주면 결과를 제대로 찾을 수 있겠죵?

 

참고로 이 방법을 이용하면 시간은 각각의 배열을 한 번씩만 훑어주면 되기 때문에

N^2 * 2의 시간이 걸린다.

 

그러면 총 시간복잡도는 입력+정렬+탐색

=O(4n + n^2 log(n^2) + 2n^2)

n=4000을 넣어보면 1억 4천

여러가지 변수를 생각해봐도 시간은 충분할 것 같다.

 

이 순서대로 코드를 정렬해주면 정답 완성!

<소스코드>

import java.io.*;
import java.util.*;

public class Main {

	public static void main(String[] args) throws IOException {
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));

		// 입력부
		int n = Integer.parseInt(br.readLine());
		int[][] arr = new int[n][4];
		for (int i = 0; i < n; i++) {
			StringTokenizer st = new StringTokenizer(br.readLine());
			for (int j = 0; j < 4; j++) {
				arr[i][j] = Integer.parseInt(st.nextToken());
			}
		}

		// 오름차순
		int[] ab = new int[n * n];
		// 내림차순
		int[] cd = new int[n * n];

		for (int i = 0; i < n; i++) {
			for (int j = 0; j < n; j++) {
				// a랑 b의 합
				ab[n * i + j] = arr[i][0] + arr[j][1];
				// c랑 d의 합
				cd[n * i + j] = arr[i][2] + arr[j][3];
			}
		}

		Arrays.sort(ab);
		Arrays.sort(cd);

		long cnt = 0;
		int l_cursor = 0;
		int r_cursor = n * n - 1;

		while (l_cursor < n * n && r_cursor >= 0) {
			int abVal = ab[l_cursor];
			int cdVal = cd[r_cursor];
			int sum = abVal + cdVal;

			if (sum == 0) { // 정답을 찾았을 경우
				long abcnt = 0;
				long cdcnt = 0;
				// 중복된 값이 없을때까지 탐색
				while (l_cursor < n * n && ab[l_cursor] == abVal) {
					l_cursor++;
					abcnt++;
				}
				while (r_cursor >= 0 && cd[r_cursor] == cdVal) {
					r_cursor--;
					cdcnt++;
				}
				cnt += (abcnt * cdcnt);
			} else if (sum > 0) { // 0보다 클 경우 cd값을 줄인다
				r_cursor--;
			} else { // 0보다 작을 경우 ab값을 높인다
				l_cursor++;
			}
		}
		System.out.println(cnt);
	}
}

 

+참고로 처음에는 우선순위 큐를 사용하여 문제를 풀려고 했었다.

배열의 정렬이 필요한 경우에 우선순위 큐를 쓰는게 이득이라고 생각되서,,,

하지만 크나큰 오산이었다.

우선순위 큐를 사용할 경우 한 번 삽입/삭제가 일어날 때 마다 nlog(n)의 시간이 소요되지만

(게다가 데이터의 수가 4000^2개)

정렬을 할 경우에는 한 번만 정렬해주면 되기 때문에 이 문제의 경우에는 배열을 쓰는게 시간적으로 이득이었다는 사실,,

 

+참고로 적어보는 우선순위 큐 코드(오답임)

import java.io.*;
import java.util.*;

public class Main {

	public static void main(String[] args) throws IOException {
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));

		// 입력부
		int n = Integer.parseInt(br.readLine());
		int[][] arr = new int[n][4];
		for (int i = 0; i < n; i++) {
			StringTokenizer st = new StringTokenizer(br.readLine());
			for (int j = 0; j < 4; j++) {
				arr[i][j] = Integer.parseInt(st.nextToken());
			}
		}

		// 오름차순
		PriorityQueue<Integer> ab = new PriorityQueue<Integer>();
		// 내림차순
		PriorityQueue<Integer> cd = new PriorityQueue<Integer>(Collections.reverseOrder());

		for (int i = 0; i < n; i++) {
			for (int j = 0; j < n; j++) {
				// a랑 b의 합
				ab.add(arr[i][0] + arr[j][1]);
				// c랑 d의 합
				cd.add(arr[i][2] + arr[j][3]);
			}
		}

		long cnt = 0;
		while (!ab.isEmpty() && !cd.isEmpty()) {
			int sum = ab.peek() + cd.peek();
			if (sum == 0) { // 정답을 찾았을 경우
				long abcnt = 1;
				long cdcnt = 1;
				// 중복된 값이 없을때까지 탐색
				while (!ab.isEmpty() && ab.poll() == ab.peek())
					abcnt++;
				while (!cd.isEmpty() && cd.poll() == cd.peek())
					cdcnt++;
				cnt += (abcnt * cdcnt);
			} else if (sum > 0) { // 0보다 클 경우 cd값을 줄인다
				cd.poll();
			} else { // 0보다 작을 경우 ab값을 높인다
				ab.poll();
			}
		}
		System.out.println(cnt);
	}
}