[백준] 2143번 두 배열의 합 - Java

문제 링크

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

문제

한 배열 A[1], A[2], …, A[n]에 대해서, 부 배열은 A[i], A[i+1], …, A[j-1], A[j] (단, 1 ≤ i ≤ j ≤ n)을 말한다. 이러한 부 배열의 합은 A[i]+…+A[j]를 의미한다. 각 원소가 정수인 두 배열 A[1], …, A[n]과 B[1], …, B[m]이 주어졌을 때, A의 부 배열의 합에 B의 부 배열의 합을 더해서 T가 되는 모든 부 배열 쌍의 개수를 구하는 프로그램을 작성하시오.

예를 들어 A = {1, 3, 1, 2}, B = {1, 3, 2}, T=5인 경우, 부 배열 쌍의 개수는 다음의 7가지 경우가 있다.

입력

첫째 줄에 T(-1,000,000,000 ≤ T ≤ 1,000,000,000)가 주어진다. 다음 줄에는 n(1 ≤ n ≤ 1,000)이 주어지고, 그 다음 줄에 n개의 정수로 A[1], …, A[n]이 주어진다. 다음 줄에는 m(1 ≤ m ≤ 1,000)이 주어지고, 그 다음 줄에 m개의 정수로 B[1], …, B[m]이 주어진다. 각각의 배열 원소는 절댓값이 1,000,000을 넘지 않는 정수이다.

출력

첫째 줄에 답을 출력한다. 가능한 경우가 한 가지도 없을 경우에는 0을 출력한다.

알고리즘: 누적합, 이분 탐색

풀이

A배열의 부 배열의 합 경우의 수는 100만 보다 작다. B배열 또한 마찬가지이다.
각 배열의 모든 부 배열 합을 구할 때는 누적합을 사용해서 구한다.
모든 부 배열 합을 구했으면, A배열의 모든 부 배열 합을 돌면서 B배열의 부 배열 합과 더했을 때 T가 되는 경우의 수가 있는지 이분 탐색으로 찾아서 카운트해주면 된다. 약 100만 * log 100만이기 때문에 충분히 통과할 수 있는 풀이법이다.

또 다른 풀이로는 hashtable을 이용한 풀이법이다. hashtable에 B배열의 모든 부 배열 합을 저장한다. 이때 hashtable에 값이 없으면 1로 세팅, 있으면 원래 값에 +1을 해준다. 그리고 A배열의 모든 부 배열 합을 돌면서 hashtable에 (T - A 부 배열 합) 값이 존재하는지 체크해주고, 존재한다면 ans에 +value 해주면 된다.

소스 코드

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

public class Main {
    static long T;
    static long A[]; //누적합
    static long B[]; //누적합
    static int N,M;
    static ArrayList<Long> ps_A = new ArrayList<>(); //부분합 list
    static Hashtable<Long, Integer> ps_B = new Hashtable<>(); //부분합 list
    static long ans = 0;
    public static void main(String args[]) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        T = Long.parseLong(br.readLine());
        N = Integer.parseInt(br.readLine());
        A = new long[N];
        StringTokenizer n_st = new StringTokenizer(br.readLine());
        for(int i=0; i<N; i++) {
            if(i==0) A[i] = Long.parseLong(n_st.nextToken());
            else A[i] = A[i-1] + Long.parseLong(n_st.nextToken());
        }
        M = Integer.parseInt(br.readLine());
        B = new long[M];
        StringTokenizer m_st = new StringTokenizer(br.readLine());
        for(int i=0; i<M; i++) {
            if(i==0) B[i] = Long.parseLong(m_st.nextToken());
            else B[i] = B[i-1] + Long.parseLong(m_st.nextToken());
        }
        for(int i=0; i<N; i++) {
            ps_A.add(A[i]);
            for(int j=0; j<i; j++) {
                ps_A.add(A[i]-A[j]);
            }
        }
        for(int i=0; i<M; i++) {
            if(ps_B.get(B[i])==null) ps_B.put(B[i],1);
            else ps_B.replace(B[i], ps_B.get(B[i])+1);
            for(int j=0; j<i; j++) {
                long pv = B[i] - B[j];
                if(ps_B.get(pv)==null) ps_B.put(pv, 1);
                else ps_B.replace(pv, ps_B.get(pv)+1);
            }
        }
        for(int i=0; i<ps_A.size(); i++) {
            long v = T - ps_A.get(i);
            if(ps_B.get(v)!=null) ans += ps_B.get(v);
        }
        System.out.println(ans);
    }
}