BOJ 7453 합이 0인 네 정수

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

문제

정수로 이루어진 크기가 같은 배열 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에 포함되는 정수가 공백으로 구분되어져서 주어진다. 배열에 들어있는 정수의 절댓값은 최대 $2^{28}$이다.

---

출력

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

---

예제 입력 1

6
-45 22 42 -16
-41 -27 56 30
-36 53 -37 77
-36 30 -75 -46
26 -38 -10 62
-32 -54 -6 45

예제 출력 1

5

---

코드

시간 제한도 12초이고, 메모리도 1024MB인 이런 문제는 뭔가 딱 어떤 개념으로 해결해야 하는지 단번에 감이 오지 않는 것 같다.

우선 DP나 그리디로 푸는 방식이 떠오르지 않아서, 다른 방법을 생각해봤다.

합이 0이 되는 경우의 수를 구하는 문제들은 투 포인터를 사용하는 문제가 많기 때문에, 주어진 4가지의 배열들을 두 개씩 짝지어 각 원소들을 더한 것들을 원소로 가지는 두 개의 배열로 나눈 뒤, 투 포인터로 0이 되는 경우를 찾아보는 방법을 떠올렸다.

풀이 방법은 다음과 같다.

1. A와 B, C와 D 배열을 짝지어 각 원소들의 합을 원소로 가지는 두 개의 배열로 만든다.
2. 만든 두 배열을 정렬시킨다.
3. 투 포인터로 두 배열의 양쪽 끝부터 차근차근 0이 되는 경우를 탐색하며 쌍의 개수를 센다.

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.*;

public class Main {

    static int[] sum1;
    static int[] sum2;

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        int N = Integer.parseInt(br.readLine());
        int[] A = new int[N];
        int[] B = new int[N];
        int[] C = new int[N];
        int[] D = new int[N];
        for (int i = 0; i < N; i++) {
            StringTokenizer st = new StringTokenizer(br.readLine());
            A[i] = Integer.parseInt(st.nextToken());
            B[i] = Integer.parseInt(st.nextToken());
            C[i] = Integer.parseInt(st.nextToken());
            D[i] = Integer.parseInt(st.nextToken());
        }

        sum1 = new int[N * N];
        sum2 = new int[N * N];
        for (int i = 0; i < N; i++) {
            for (int j = 0; j < N; j++) {
                sum1[i * N + j] = A[i] + B[j];
                sum2[i * N + j] = C[i] + D[j];
            }
        }

        Arrays.sort(sum1);
        Arrays.sort(sum2);

        long answer = countCases();
        System.out.println(answer);
    }

    static long countCases() {
        long count = 0;
        int index1 = 0;
        int index2 = sum2.length - 1;

        while (index1 < sum1.length && index2 >= 0) {
            int sum = sum1[index1] + sum2[index2];
            if(sum > 0) {
                index2--;
            } else if (sum < 0) {
                index1++;
            } else {
                int sameIndex1 = index1 + 1;
                int sameIndex2 = index2 - 1;
                while (sameIndex1 < sum1.length && sum1[index1] == sum1[sameIndex1]) {
                    sameIndex1++;
                }
                while (sameIndex2 >= 0 && sum2[index2] == sum2[sameIndex2]) {
                    sameIndex2--;
                }

                count += (long) (sameIndex1 - index1) * (index2 - sameIndex2);
                index1 = sameIndex1;
                index2 = sameIndex2;
            }
        }

        return count;
    }
}