[JAVA] 백준 (골드3) 2143번 두 배열의 합

링크

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

---

입력 예제

5
4
1 3 1 2
3
1 3 2

---

출력 예제

7

---

풀이

A 부분 배열의 모든 경우의 수를 구하는 건 $O(n^2)$, B 부분 배열의 모든 경우의 수를 구하는 건 $O(m^2)$이다. 여기다가 각 부분 배열의 합을 구하는 경우의 수는 $O(n^2 × m^2)$가 되어 비효율적이다.

따라서 모든 부분 배열의 경우의 수를 구한 뒤, 정렬하여 투 포인터를 이용해서 합이 T가 될 때를 구하면 시간 복잡도는 $O(N^2logN+M^2logM)$로 최적화할 수 있다.

A의 구간 합 리스트는 오름차순으로, B의 구간 합 리스트는 내림차순으로 정렬한다.

//A 구간 합의 모든 경우의 수
List<Integer> listA = new ArrayList<>();
for (int i = n; i > 0; i--) {
    for (int j = i; j > 0; j--) {
        int x = sumA[i] - sumA[j - 1];
        listA.add(x);
    }
}

//B 구간 합의 모든 경우의 수
List<Integer> listB = new ArrayList<>();
for (int i = m; i > 0; i--) {
    for (int j = i; j > 0; j--) {
        int y = sumB[i] - sumB[j - 1];
        listB.add(y);
    }
}

listA.sort(Comparator.naturalOrder());
listB.sort(Comparator.reverseOrder());

그리고 투 포인터를 이용하여 각 구간 합의 합이 T가 될 때 가능한 경우의 수를 카운트한다.

//투 포인터
int s = 0;
int e = 0;
while (s < listA.size() && e < listB.size()) {
    int a = listA.get(s);  //A 부 배열의 합
    int b = listB.get(e);  //B 부 배열의 합
    int sum = a + b;  //부 배열의 합
    
    if (sum == T) {
        long countA = 0;  //중복된 수 카운트
        while (s < listA.size() && listA.get(s) == a) {
            countA++;
            s++;  //다음 값으로 넘어감
        }
        
        long countB = 0;
        while (e < listB.size() && listB.get(e) == b) {
            countB++;
            e++;
        }
        
        count += countA * countB;  //경우의 수
    } else if (sum < T) {
        s++;
    } else {
        e++;
    }
}

전체 코드

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;
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 = Integer.parseInt(br.readLine());

        //A배열 초기화
        int n = Integer.parseInt(br.readLine());
        int[] A = new int[n + 1];
        StringTokenizer st = new StringTokenizer(br.readLine());
        for (int i = 1; i <= n; i++) {
            A[i] = Integer.parseInt(st.nextToken());
        }

        //A배열 구간 합
        int[] sumA = new int[n + 1];
        for (int i = 1; i <= n; i++) {
            sumA[i] = sumA[i - 1] + A[i];
        }

        //B배열 초기화
        int m = Integer.parseInt(br.readLine());
        int[] B = new int[m + 1];
        st = new StringTokenizer(br.readLine());
        for (int i = 1; i <= m; i++) {
            B[i] = Integer.parseInt(st.nextToken());
        }

        //B배열 구간 합
        int[] sumB = new int[m + 1];
        for (int i = 1; i <= m; i++) {
            sumB[i] = sumB[i - 1] + B[i];
        }

        //A 구간 합의 모든 경우의 수
        List<Integer> listA = new ArrayList<>();
        for (int i = n; i > 0; i--) {
            for (int j = i; j > 0; j--) {
                int x = sumA[i] - sumA[j - 1];
                listA.add(x);
            }
        }

        //B 구간 합의 모든 경우의 수
        List<Integer> listB = new ArrayList<>();
        for (int i = m; i > 0; i--) {
            for (int j = i; j > 0; j--) {
                int y = sumB[i] - sumB[j - 1];
                listB.add(y);
            }
        }

        listA.sort(Comparator.naturalOrder());
        listB.sort(Comparator.reverseOrder());

        long count = 0;

        //투 포인터
        int s = 0;
        int e = 0;
        while (s < listA.size() && e < listB.size()) {
            int a = listA.get(s);  //A 부 배열의 합
            int b = listB.get(e);  //B 부 배열의 합
            int sum = a + b;  //부 배열의 합

            if (sum == T) {
                long countA = 0;  //중복된 수 카운트
                while (s < listA.size() && listA.get(s) == a) {
                    countA++;
                    s++;  //다음 값으로 넘어감
                }

                long countB = 0;
                while (e < listB.size() && listB.get(e) == b) {
                    countB++;
                    e++;
                }

                count += countA * countB;  //경우의 수
            } else if (sum < T) {
                s++;
            } else {
                e++;
            }
        }

        System.out.println(count);
    }
}