[BOJ 2143] 두 배열의 합 (Java)

BOJ 2143 두 배열의 합

난이도 : 골드 3

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

2143번: 두 배열의 합

 

전략

처음에는 prefix sum을 구해놓고 포인터를 움직여가는 방식은 어떨까 생각해봤는데 배열이 두개라서 상당히 복잡할 것 같았다.

따라서 부배열 합들을 모두 저장해놓고 탐색하는 것이 좋겠다고 생각했다.

길이가 N인 배열에서 모든 부배열의 합을 구하는 것은 O(N^2)이 걸린다. 하지만 문제에서 N이 최대 1000이므로 충분하다.

또한 길이가 M인 배열에 대해서도 O(M^2)만큼의 전처리가 같은 방법으로 필요하다.

 

그 다음 N^2개의 저장된 합에 대해 길이가 M인 배열의 부배열 합 중 더해서 T가 될 수 있는 것을 찾아야 한다.

생각한 방법은 이분 탐색 또는 HashMap을 이용하는 것이다.

HashMap은 해시 충돌로 최악의 경우에 O(M)이 걸릴 수 있을 것이다. 만약 O(N^2 * M)까지 간다면 시간 안에 해결할 수 없을 것이라 생각해서 이분 탐색으로 결정했다.

 

따라서 길이가 M인 부배열의 합들을 미리 정렬해놓고, 이분탐색을 진행한다.

부배열 합들의 크기는 M^2이므로 정렬에는 M^2*log(M*M)이므로 O(M^2*logM)

그렇다면 최종적인 시간복잡도는 O(N^2 * logM + M^2 * logM)이 된다.

 

시행착오

        for (int sum : subSumA) {
            if (sum >= T)
                break;
            if (binSearch(subSumB, T - sum) != -1)
                cnt++;
        }

다음과 같이 이분탐색을 진행해서 일치하는 것이 있으면 카운트를 1만 증가시켰다.

만약 일치하는 것이 여러개 있다면 어떻게 될까?

예를 들어 T가 5고 subSumA 배열에서 부분합 3이 선택되었을 때

subSumB에서 부분합이 2가 되는 부배열이 여러 개 있다면 그만큼 세어야 하지만, 위의 코드로는 무조건 1개로 세는 것이다.

따라서 이분 탐색을 두 번하여 upper bound와 lower bound를 각각 구하고 그 수만큼 더해줘야한다.!

 

코드

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.Collections;
import java.util.StringTokenizer;

public class Main {

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        int T, N, M;
        T = Integer.parseInt(br.readLine());
        N = Integer.parseInt(br.readLine());
        int[] arrayA = new int[N];
        StringTokenizer st = new StringTokenizer(br.readLine());
        for (int i=0; i<N; i++)
            arrayA[i] = Integer.parseInt(st.nextToken());

        M = Integer.parseInt(br.readLine());
        int[] arrayB = new int[M];
        st = new StringTokenizer(br.readLine());
        for (int i=0; i<M; i++)
            arrayB[i] = Integer.parseInt(st.nextToken());

        ArrayList<Integer> subSumA = getSubSumList(arrayA);
        ArrayList<Integer> subSumB = getSubSumList(arrayB);
        long cnt = 0;

        Collections.sort(subSumB);

        for (int sum : subSumA) {
            int tar = T - sum;
            cnt += upperBound(subSumB, tar) - lowerBound(subSumB, tar);
        }
        System.out.println(cnt);
    }

    static int upperBound(final ArrayList<Integer> list, int find) {
        int left = 0;
        int right = list.size();
        int mid;

        while (left < right) {
            mid = (left + right) / 2;
            if (list.get(mid) <= find)
                left = mid + 1;
            else
                right = mid;
        }
        return right;
    }

    static int lowerBound(final ArrayList<Integer> list, int find) {
        int left = 0;
        int right = list.size();
        int mid;

        while (left < right) {
            mid = (left + right) / 2;
            if (list.get(mid) < find)
                left = mid + 1;
            else
                right = mid;
        }
        return right;
    }

    static ArrayList<Integer> getSubSumList(final int[] arr) {
        ArrayList<Integer> list = new ArrayList<>();

        for (int i=0; i<arr.length; i++) {
            int sum = 0;
            for (int j=i; j<arr.length; j++) {
                sum += arr[j];
                list.add(sum);
            }
        }
        return list;
    }
}

 

깨달은 점

이분탐색으로 upperBound, lowerBound를 짤 때 헷갈리기 쉽다고 느꼈다.

특히 right 시작 위치를 습관적으로 list.size() - 1로 정한다면

리스트 크기가 1인 경우에 upperBound가 틀리는 것을 볼 수 있다. 앞으론 이런 부분 더 신경써서 할 수 있을 것 같다.