백준 7453 합이 0인 네 정수 c++, Kotlin (투 포인터)

문제 출처 : https://www.acmicpc.net/problem/7453

7453번: 합이 0인 네 정수

문제

정수로 이루어진 크기가 같은 배열 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이 되는 쌍의 개수를 출력한다.

알고리즘 분류

• 정렬
• 이분 탐색
• 두 포인터
• 중간에서 만나기

풀이

크기가 4000인 배열이 4개이기 때문에 완전탐색으로는 O(N^4)로 시간초과다.

딱히 풀이가 생각나지 않아서 구글링하였고, 투 포인터 공부의 필요성을 느꼈다.

사실 투 포인터란 알고리즘이 있는지도 모른 채 그냥 구현이다 생각하면서 사용해왔다.

이번 기회에 투 포인터 알고리즘에 대한 개념을 정립했다.

최근 어느 알고리즘 대회에 투 포인터 문제가 나오기도 했고,

일반적인 투 포인터 알고리즘의 시간 복잡도는 O(N)으로

어느 리스트에서 구간 합이 m인 구간의 개수를 구하는 문제에서,

2중 포문을 사용해야 하는 완전탐색에 비해 엄청난 효율을 보여준다.

 

각설하고, 투 포인터를 이용한 풀이는 다음과 같다.

1. A,B배열의 각 요소들을 더하여 배열 1을 만들고 C,D배열의 각 요소들을 더하여 배열 2를 만든다.

2. 배열 1과 배열 2를 오름차순으로 정렬한다.

3. 배열 1의 요소는 a+b이고, 배열 2의 요소는 c+d이기 때문에 배열 1의 요소와 배열 2의 요소를 더한 값이 0인 것을 찾
   으면 된다.

4. 여기서 투 포인터를 이용하게 되는데 p1은 배열 1의 가장 왼쪽을 가리키고, p2은 배열 2의 가장 오른쪽을 가리킨다.

5. 두 배열을 모두 오름차순 해놨기에, p1이 커질수록 값이 커지고, p2가 작아질수록 값이 작아진다.

6. arr1[p1]+arr2[p2]의 값이 sum이라 할 때 sum이 0보다 작으면 p1을 증가시킨다. (sum이 커진다)

7. arr1[p1]+arr2[p2]의 값이 sum이라 할 때 sum이 0보다 크면 p2를 감소시킨다. (sum이 작아진다)

8. arr1[p1]+ arr2[p2]의 값이 sum이라 할 때 sum이 0이라면, arr1[p1]과 값이 같은 원소의 개수와 arr2[p2]와 값이 같은
   원소의 개수를 곱한다.

ex) 

배열 1 { -4, -4, -4 ,-2 , -1, 2 ,6}

배열 2 { -3, -1, -2 ,0, 2, 4, 4}

p1이 0이고 p2가 6일 때,

-4 +4로 sum 은 0이다.

이때, arr1에서 -4는 3개, arr2에서 4는 2개이기 때문에,

-4 +4로 sum이 0이 되는 경우는 총 6가지이다.

이후 p1은 -4 이후인 3으로 이동하고, p2는 4 이후인 2로 이동한다.

 

9. 위의 6번부터 과정을 p1 혹은 p2가 배열을 이탈할 때까지 반복한다.

코드(c++)

#include <iostream>
#include <algorithm>
#include <vector>
#define MAX 4000

using namespace std;
long long arr1[MAX*MAX], arr2[MAX*MAX];
int main() {
	ios::sync_with_stdio(false);
	cin.tie(0);
	cout.tie(0);
	int arr[4][MAX];
	long long answer = 0;
	int n;
	cin >> n;
	for (int i = 0; i < n; i++) {
		cin >> arr[0][i] >> arr[1][i] >> arr[2][i] >> arr[3][i];
	}
	long long idx = 0;
	for (int i = 0; i < n; i++) {
		for (int j = 0; j < n; j++) {
			arr1[idx] = arr[0][i] + arr[1][j];
			arr2[idx] = arr[2][i] + arr[3][j];
			idx++;
		}
	}
	sort(arr1, arr1 + idx);
	sort(arr2, arr2 + idx);

	int p1 = 0, p2 = idx- 1;

	while (p1 < idx && p2 >= 0) {
		if (arr1[p1] + arr2[p2] == 0) {
			long long cnt1 = 0;
			int i = p1;
			while (i < idx && arr1[i] == arr1[p1]) {
				cnt1++;
				i++;
			}
			p1 = i;

			long long cnt2 = 0;
			i = p2;
			while (i >= 0 && arr2[i] == arr2[p2]) {
				cnt2++;
				i--;
			}
			p2 = i;
			answer += cnt1 * cnt2;
		}
		else if (arr1[p1] + arr2[p2] < 0) {
			p1++;
		}
		else {
			p2--;
		}
	}
	cout << answer;


	return 0;
}

코드(Kotlin)

import java.util.StringTokenizer
fun main()= with(System.out.bufferedWriter()){
    val br = System.`in`.bufferedReader()
    val n = Integer.parseInt(br.readLine())
    val arr = Array<LongArray>(4){LongArray(n)}
    val arr1 = LongArray(n*n)
    val arr2 = LongArray(n*n)
    var ans = 0L
    for(i in 0 until n){
        val tk = StringTokenizer(br.readLine())
        for(j in 0 until 4){
            arr[j][i] = tk.nextToken().toLong()
        }
    }
    for(i in 0 until n){
        for(j in 0 until n){
            arr1[i*n+j] = arr[0][i] + arr[1][j]
            arr2[i*n+j] = arr[2][i] + arr[3][j]
        }
    }

    arr1.sort()
    arr2.sort()

    var p1 = 0
    var p2 = arr2.size-1

    while(p1<arr1.size && p2>=0){
        if(arr1[p1]+arr2[p2]==0L){
            var i = p1
            var cnt1=0L
            while(i <arr1.size && arr1[i]==arr1[p1]){
                cnt1++
                i++
            }
            p1=i
            i =p2
            var cnt2 =0L
            while(i>=0&&arr2[i]==arr2[p2]){
                cnt2++
                i--
            }
            ans+=cnt1*cnt2
            p2=i
        }
        else if(arr1[p1]+arr2[p2]<0){
            p1++
        }
        else{
            p2--
        }
    }
    write("$ans")
    close()
}