[백준] 2143번 - 두 배열의 합 (Gold 3)

문제 링크

백준 2143번 - 두 배열의 합 (Gold 3)

문제 설명

두 개의 정수로 이루어진 수열 A, B가 주어져있다.
주어진 정수 T에 대해, T = A의 연속된 부분수열의 합 + B의 연속된 부분수열의 합이 되는 경우의 수를 구하시오.
같은 수열이어도 시작과 끝나는 인덱스가 다르다면 다른 부분수열로 취급한다. 예를 들어 A = {1, 2, 1, 2}라면 부분수열 {A[1], A[2]} = {A[3], A[4]} = {1, 2}지만 두 부분수열은 다른 것으로 센다.

정답 코드 및 설명

연속된 부분수열의 합은 누적합으로 쉽게 구할 수 있다.
주어진 조건을 이항하면 A의 연속된 부분수열의 합 = T - B의 연속된 부분수열의 합인 경우의 수이다.
A의 연속된 부분수열의 합을 전부 구해 key가 부분수열의 합, value가 합이 key인 부분수열의 개수인 맵을 만들자. 전체 경우의 수는 B의 각 연속된 부분수열마다 맵에서 key가 T - B의 연속된 부분수열의 합인 value(해당하는 key가 없는 경우는 0)를 꺼내 모두 더한 값이다.

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.HashMap;
import java.util.StringTokenizer;

public class BOJ2143 {
    int t;
    int[] a, b;
    HashMap<Integer, Integer> sumA;
    long count = 0;

    void input() throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        t = Integer.parseInt(br.readLine());
        a = setArr(br);
        sumA = setSum(a);
        b = setArr(br);
    }

    int[] setArr(BufferedReader br) throws IOException {
        int length = Integer.parseInt(br.readLine()) + 1;
        int[] a = new int[length];
        StringTokenizer st = new StringTokenizer(br.readLine());
        for (int i = 1; i < length; i++) {
            a[i] = a[i - 1] + Integer.parseInt(st.nextToken());
        }
        return a;
    }

    HashMap<Integer, Integer> setSum(int[] a) {
        HashMap<Integer, Integer> sum = new HashMap<>();
        for (int i = 0; i < a.length; i++) {
            for (int j = 0; j < i; j++) {
                sum.put(a[i] - a[j], sum.getOrDefault(a[i] - a[j], 0) + 1);
            }
        }
        return sum;
    }

    void solution() throws IOException {
        input();
        for (int i = 0; i < b.length; i++) {
            for (int j = 0; j < i; j++) {
                count += sumA.getOrDefault(t - (b[i] - b[j]), 0);
            }
        }
        System.out.println(count);
    }

    public static void main(String[] args) throws IOException {
        new BOJ2143().solution();
    }
}