Article Image
Article Image
read

문제

한 배열 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(=5) = A[1] + B[1] + B[2]
      = A[1] + A[2] + B[1]
      = A[2] + B[3]
      = A[2] + A[3] + B[1]
      = A[3] + B[1] + B[2]
      = A[3] + A[4] + B[3]
      = A[4] + B[2] 

입력

첫째 줄에 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을 출력한다.

예제 입력 1

5
4
1 3 1 2
3
1 3 2

예제 출력 1

7

풀이

이 문제는 누적합과 이분탐색을 이용해서 풀 수 .

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.Collections;

public class Main {

    public static void main(String[] args) throws Exception{
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        int T = Integer.parseInt(br.readLine());
        int N = Integer.parseInt(br.readLine());
        String[] input = br.readLine().split(" ");

        int[] arr = new int[N];
        for(int i=0; i<N; i++) {
            arr[i] = Integer.parseInt(input[i]);
        }

        ArrayList<Long> A = new ArrayList<>();
        for(int i=0; i<N; i++) {
            long sum = 0;
            for(int j=i; j<N; j++) {
                sum += arr[j];
                A.add(sum);
            }
        }

        int M = Integer.parseInt(br.readLine());
        input = br.readLine().split(" ");

        arr = new int[M];
        for(int i=0; i<M; i++) {
            arr[i] = Integer.parseInt(input[i]);
        }

        ArrayList<Long> B = new ArrayList<>();
        for(int i=0; i<M; i++) {
            long sum = 0;
            for(int j=i; j<M; j++) {
                sum += arr[j];
                B.add(sum);
            }
        }

        Collections.sort(A);
        Collections.sort(B);

        int left = 0;
        int right = B.size()-1;
        long sum = 0;

        while(left<A.size() && right>=0) {
            long a = A.get(left);
            long b = B.get(right);

            if(a+b == T) {
                int cntA = 0;
                int cntB = 0;

                while(left<A.size() && a==A.get(left)) {
                    left++;
                    cntA++;
                }

                while(right>=0 && b==B.get(right)) {
                    right--;
                    cntB++;
                }

                sum += (long)cntA*cntB;
            }

            else if(a+b > T) {
                right--;
            }

            else {
                left++;
            }
        }

        System.out.println(sum);
    }
}
Blog Logo

pss407


Published

Image

pss's GitHub 블로그

Welcome my blog!

Back to Overview