[백준] 7453: 합이 0인 네 정수 (이분탐색, upper bound & lower bound) - JAVA

class 5+

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에 포함되는 정수가 공백으로 구분되어져서 주어진다. 배열에 들어있는 정수의 절댓값은 최대 228이다.

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

 

 

출력

5

 

# 필요개념

배열의 개수가 4개로 정해져있지만, 배열의 최대 크기가 4000 이고, 시간은 12초 뿐이기 때문에 시간복잡도 관리에 유의해야 한다. 또, 배열에 들어가는 정수의 절댓값도 큰 편이므로, 이 점도 유의한다.

 

4중 for문을 돌리는 건 말도 안 되는 풀이이고, 어떻게 풀까 하다가 A + B 배열, C + D 배열을 만들어두기로 했다. A + B에 5라는 값이 있으면, C + D 배열에 -5가 몇 개인지 확인하면 되기 때문이다. 

 

이를 확인할 때, for문을 돌려 앞에서부터 확인하면 시간초과가 발생한다. 그래서 이분 탐색을 사용했다.

 

✅ Upper & Lower

단순히 유무가 아니라 조건에 맞는 수가 몇 개 들어있는지를 알아야 하기 때문에, upper & Lower 방식을 사용했다.

이분탐색을 따르되, 원하는 수가 가장 처음 등장하는 지점 (Lower) 과 원하는 수보다 큰 수가 처음 나오는 지점 (Upper) 의 인덱스를 찾아내면 된다. 그렇게 하면 upper Idx - lower Idx 가 원하는 수의 개수인 게 된다.

 

 

mid를 기준으로 범위를 좁혀나가면 시간을 단축할 수 있다 !

이때 A + B 배열과 C + D 배열은 int의 범위를 벗어날 수 있기 때문에 long 타입으로 선언해주었다.

 

---

# Code

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

public class Main {
    static int N;
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        N = Integer.parseInt(br.readLine());
        int[][] inputs = new int[N][4];
        for (int i = 0 ; i < N ; i++) {
            StringTokenizer st = new StringTokenizer(br.readLine());
            inputs[i][0] = Integer.parseInt(st.nextToken());
            inputs[i][1] = Integer.parseInt(st.nextToken());
            inputs[i][2] = Integer.parseInt(st.nextToken());
            inputs[i][3] = Integer.parseInt(st.nextToken());
        }
        /*
         * 단순히 contains 쓰면 조건에 맞는 합의 개수를 알 수 없음 -> upperBound & LowerBound 사용
         * */
        long[] sumAB = new long[N * N];
        long[] sumCD = new long[N * N]; // 정렬해서 upper & LowerBound
        int idx = 0;
        for (int i = 0 ; i < N ; i++) {
            for (int j = 0 ; j < N ; j++) {
                sumAB[idx] = inputs[i][0] + inputs[j][1];
                sumCD[idx] = inputs[i][2] + inputs[j][3];
                idx++;
            }
        }
        Arrays.sort(sumCD);

        long ans = 0;
        for (int i = 0 ; i < sumAB.length ; i++) {
            int upperBoundIdx = upperBound(-1 * sumAB[i], sumCD);
            int lowerBoundIdx = lowerBound(-1 * sumAB[i], sumCD);
            if (lowerBoundIdx <= upperBoundIdx) ans += (upperBoundIdx - lowerBoundIdx);
        }
        System.out.println(ans);
    }

    private static int lowerBound(long goal, long[] sumCD) {
        int low = 0;
        int high = sumCD.length;
        while (low < high) {
            int mid = (low + high) / 2;
            if (sumCD[mid] >= goal) high = mid;
            else low = mid + 1;
        }
        return low;
    }

    private static int upperBound(long goal, long[] sumCD) {
        int low = 0;
        int high = sumCD.length;
        while (low < high) {
            int mid = (low + high) / 2;
            if (sumCD[mid] <= goal) low = mid + 1;
            else high = mid;
        }
        return low;
    }
}

 

 

# 결과