b.2143 두 배열의 합

 

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

 

 

 

이 문제는 두 배열의 부분합 중 합이 t가 되는 모든 쌍의 개수를 구하는 문제다.
먼저 두 배열 A와 B 각각의 부분합을 계산한다.

각 배열의 모든 부분합은 시작점과 끝점을 조합하여 구간의 합을 계산하는 방식으로 얻어진다.
이 과정은 calculateSubarraySums 메서드에서 수행된다.

부분합 배열 aSum과 bSum은 투 포인터 알고리즘을 사용하기 위해 정렬된다.
이후 calculatePairCount 메서드에서 두 배열을 탐색하며 합이 t인 쌍의 개수를 계산한다.

이 과정에서는 각 부분합의 포인터를 이동시키며 두 값의 합이 t와 같을 경우,
현재 값과 동일한 원소의 개수를 세어 곱한 값을 결과에 추가한다.
합이 t보다 크면 bSum 포인터를 줄이고,
작으면 aSum 포인터를 늘리며 조건을 만족하는 쌍을 찾아간다.

결과적으로 모든 쌍을 효율적으로 탐색하며 조건을 만족하는 경우의 수를 구한다.

 

 

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

public class Main {

    static int an, bn, t;
    static int[] A, B;
    static int[] aSum, bSum;

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));

        t = Integer.parseInt(br.readLine());
        an = Integer.parseInt(br.readLine());
        A = Arrays.stream(br.readLine().split(" ")).mapToInt(Integer::parseInt).toArray();
        bn = Integer.parseInt(br.readLine());
        B = Arrays.stream(br.readLine().split(" ")).mapToInt(Integer::parseInt).toArray();

        aSum = calculateSubarraySums(A, an);
        bSum = calculateSubarraySums(B, bn);

        Arrays.sort(aSum);
        Arrays.sort(bSum);

        long result = calculatePairCount(aSum, bSum, t);

        System.out.println(result);
    }

    static int[] calculateSubarraySums(int[] array, int size) {
        int[] subarraySums = new int[size * (size + 1) / 2];
        int index = 0;

        for (int i = 0; i < size; i++) {
            int sum = 0;
            for (int j = i; j < size; j++) {
                sum += array[j];
                subarraySums[index++] = sum;
            }
        }

        return subarraySums;
    }

    static long calculatePairCount(int[] aSum, int[] bSum, int target) {
        int left = 0, right = bSum.length - 1;
        long count = 0;

        while (left < aSum.length && right >= 0) {
            long asv = aSum[left], bsv = bSum[right];
            long sum = asv + bsv;

            if (sum == target) {
                long ac = 0, bc = 0;

                while (left < aSum.length && aSum[left] == asv) {
                    left++;
                    ac++;
                }

                while (right >= 0 && bSum[right] == bsv) {
                    right--;
                    bc++;
                }

                count += ac * bc;

            } else if (sum > target) {
                right--;
            } else {
                left++;
            }
        }

        return count;
    }
}