[BOJ-2143] 두 배열의 합(JAVA)

백준 2143 두 배열의 합

 

2143번: 두 배열의 합

 

문제 설명

 

 - 한 배열이 A[1], A[2], ..., A[n] 라고 가정할 때 부 배열은 A[i], A[i+1], ..., A[j]이고, ( 1 < = i <= j <= n )

   부 배열의 합을 A[i] + ... + A[j]라고 하자.

 

 - 두 배열 A[1], ... A[n]과 B[1], ..., B[m]이 주어졌을 때,

   A의 부 배열의 합과 B의 부 배열의 합을 더해 T가 되는 부 배열의 쌍을 구하라.(가능한 경우가 없는 경우 0을 출력하라)

 

 - 예를 들어, 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 을 넘지 않는 정수이다.

 

---

 

 

예제

 

---

 

 

문제 분석

 

 - 이 문제는 두 배열의 연속된 배열 합이 T인 경우의 수를 구하는 문제이다.

 

 - 즉, 모든 연속합의 경우의 수를 구하고 그 합이 T인 경우를 모두 구하면 된다.

 

 - 연속합의 모든 경우를 구하는 방법은 Prefix Sum을 통해 구하면 된다.

 

 - 배열의 크기가 n일 때(arr[1], arr[2], ... ,arr[n]) sum[i] = arr[1] + arr[2] + ... + arr[i] 라고 하자.

 

 - 이렇게 sum 배열을 구하고 나면, sum[i] - sum[j-1] (i > j) = j ~ i 까지의 연속 합을 구할 수 있다.

 - 이렇게 모든 연속된 구간합의 경우를 구할 수 있다.

 

 - 이 경우, 배열의 크기가 n일 때 O($n^2$)의 시간 복잡도를 가진다.

 

 - 주어진 두 배열 A와 B의 모든 연속 구간합을 구한다.

 

 - 그 다음, 배열 A의 모든 구간합 중 하나를 고른다. 그 이후에 배열 B의 모든 구간합에서 A에서 선택한 값과

   더하여 T가 되는 경우의 수를 구한다.

 

 - 이때 배열 B의 모든 구간합에서 값을 찾는 방법은 TreeMap을 이용하면 된다.

 

 - 즉, 시간 복잡도는 각 배열의 모든 부분 배열을 구하는 시간 복잡도 $N^2$ + $N^2$ + $N^2$logN 이다

---

 

 

소스 코드

 

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

// https://www.acmicpc.net/problem/2143
// 두 배열의 합
public class Main {

    public static long sumA[] = new long[1001];
    public static long sumB[] = new long[1001];

    public static List<Long> sumListA = new ArrayList<>();
    public static Map<Long,Long> sumMapB = new TreeMap<>();

    public static void main(String[] args) throws IOException {

        BufferedReader bf = new BufferedReader(new InputStreamReader(System.in));

        int T = Integer.parseInt(bf.readLine());

        int N = Integer.parseInt(bf.readLine());

        long result=0;

        StringTokenizer st = new StringTokenizer(bf.readLine());

        for(int i=0;i<N;i++){
            int num = Integer.parseInt(st.nextToken());
            for(int j=i;j<N;j++){
                sumA[j]+=num;
            }
        }

        int M = Integer.parseInt(bf.readLine());

        st = new StringTokenizer(bf.readLine());

        for(int i=0;i<M;i++){
            int num = Integer.parseInt(st.nextToken());
            for(int j=i;j<M;j++){
                sumB[j]+=num;
            }
        }

        for(int i=0;i<N;i++){
            sumListA.add(sumA[i]);
            for(int j=i+1;j<N;j++){
                sumListA.add(sumA[j]-sumA[i]);
            }
        }

        for(int i=0;i<M;i++){
            long cnt = sumMapB.containsKey(sumB[i]) ? sumMapB.get(sumB[i]) : 0;
            sumMapB.put(sumB[i], cnt+1);
            for(int j=i+1;j<M;j++){
                cnt = sumMapB.containsKey(sumB[j]-sumB[i]) ? sumMapB.get(sumB[j]-sumB[i]) : 0;
                sumMapB.put(sumB[j]-sumB[i], cnt+1);
            }
        }

        Collections.sort(sumListA);

        for(long a : sumListA){
            result += sumMapB.containsKey(T-a) ? sumMapB.get(T-a) : 0;
        }

        System.out.println(result);

    }

}