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

1.  문제 링크

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

2.  문제

요약

• 한 배열 A[1], A[2], ..., A[n]에 대해서, 부 배열은 A[i], A[i + 1], ..., A[j - 1], A[j]를 말합니다.
• 부 배열의 합은 A[i] + A[i + 1] + ... + A[j]를 의미합니다.
• 각 원소가 정수인 두 배열 A[1], A[2], ..., A[n]과 B[1], B[2], ..., B[m]이 주어졌을 때, A의 부배열의 합에 B의 부 배열의 합을 더해서 T가 되는 모든 부 배열 쌍의 개수를 구하는 문제입니다.
• 입력: 첫 번째 줄에 -1,000,000,000보다 크거나 같고 1,000,000,000보다 작거나 같은 T가 주어지고 두 번째 줄에 1보다 크거나 같고 1,000보다 작거나 같은 n이 주어지며 세 번째 줄에 n개의 정수로 A[1], ..., A[n]이 주어지고 네 번째 줄에는 1보다 크거나 같고 1,000보다 작거나 같은 m이 주어지며 마지막 줄에는 m개의 정수로 B[1], ..., [m]이 주어집니다.
  • 각 배열 원소는 절댓값이 1,000,000을 넘지 않는 정수입니다.
• 출력: 첫 번째 답을 출력합니다. 가능한 경우가 한 가지도 없을 경우에는 0을 출력합니다.

3.  소스코드

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

public class Main {
    static int T, n, m;
    static int[] A, B;
    static void input() {
        Reader scanner = new Reader();
        T = scanner.nextInt();
        n = scanner.nextInt();
        A = new int[n];
        for(int idx = 0; idx < n; idx++) A[idx] = scanner.nextInt();
        m = scanner.nextInt();
        B = new int[m];
        for(int idx = 0; idx < m; idx++) B[idx] = scanner.nextInt();
    }

    static void solution() {
        int[] prefixSumA = prefixSum(A, n), prefixSumB = prefixSum(B, m);
        HashMap<Integer, Integer> sumA = getNumOfAllSum(prefixSumA), sumB = getNumOfAllSum(prefixSumB);
        List<Map.Entry<Integer, Integer>> listA = getSortedMap(sumA), listB = getSortedMap(sumB);
        long answer = binarySearch(listA, listB);
        System.out.println(answer);
    }

    static long binarySearch(List<Map.Entry<Integer, Integer>> listA, List<Map.Entry<Integer, Integer>> listB) {
        long answer = 0L;
        int pointerA = 0, pointerB = listB.size() - 1;
        while(pointerA < listA.size() && pointerB >= 0) {
            int sum = listA.get(pointerA).getKey() + listB.get(pointerB).getKey();
            if(sum == T) {
                long numA = (long)listA.get(pointerA).getValue();
                pointerA++;
                long numB = (long)listB.get(pointerB).getValue();
                pointerB--;
                answer += numA * numB;
            } else if(sum > T) {
                pointerB--;
            } else {
                pointerA++;
            }
        }
        return answer;
    }

    static List<Map.Entry<Integer, Integer>> getSortedMap(HashMap<Integer, Integer> sum) {
        List<Map.Entry<Integer, Integer>> list = new ArrayList<>(sum.entrySet());
        Collections.sort(list, new Comparator<Map.Entry<Integer, Integer>>() {
            @Override
            public int compare(Map.Entry<Integer, Integer> o1, Map.Entry<Integer, Integer> o2) {
               return o1.getKey() - o2.getKey();
            }
        });
        return list;
    }

    static HashMap<Integer, Integer> getNumOfAllSum(int[] prefixSum) {
        HashMap<Integer, Integer> sum;
        sum = new HashMap<>();
        for(int end = 1; end < prefixSum.length; end++) {
            for(int start = 0; start < end; start++) {
                int value = prefixSum[end] - prefixSum[start];
                sum.put(value, sum.getOrDefault(value, 0) + 1);
            }
        }
        return sum;
    }

    static int[] prefixSum(int[] org, int size) {
        int[] arr = new int[size + 1];
        for(int idx = 1; idx <= size; idx++) {
            arr[idx] = org[idx - 1];
            arr[idx] += arr[idx - 1];
        }
        return arr;
    }

    public static void main(String[] args) {
        input();
        solution();
    }

    static class Reader {
        BufferedReader br;
        StringTokenizer st;
        public Reader() {
            br = new BufferedReader(new InputStreamReader(System.in));
        }
        String next() {
            while (st == null || !st.hasMoreElements()) {
                try {
                    st = new StringTokenizer(br.readLine());
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
            return st.nextToken();
        }
        int nextInt() {
            return Integer.parseInt(next());
        }
    }
}