[백준] 2143번 두 배열의 합 (JAVA)

문제

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

2143번: 두 배열의 합

 

설명

A와 B의 임의의 부 배열 두 개의 합을 더해서 T가 되는 모든 부 배열 쌍의 개수를 구해야 한다.

먼저 A와 B의 원소를 저장할 배열인 arrA와 arrB를 만든다.

그 다음, A와 B의 모든 부 배열의 각 원소의 합을 저장할 리스트인 sumA와 sumB를 만든다.

모든 부 배열의 각 원소의 합을 저장 한 뒤, sumA는 오름차순으로, sumB는 내림차순으로 정렬한다.

List<Long> sumA = new ArrayList<>();
List<Long> sumB = new ArrayList<>();

for (int i = 0; i < n; i++) {
    long sum = arrA[i];
    sumA.add(sum);
    for (int j = i + 1; j < n; j++) {
        sum += arrA[j];
        sumA.add(sum);
    }
}

for (int i = 0; i < m; i++) {
    long sum = arrB[i];
    sumB.add(sum);
    for (int j = i + 1; j < m; j++) {
        sum += arrB[j];
        sumB.add(sum);
    }
}

Collections.sort(sumA);
Collections.sort(sumB, Comparator.reverseOrder());

두 리스트를 움직일 포인터인 ptA와 ptB를 만들고 각 리스트의 앞에 위치시킨다.

현재 sumA는 오름차순으로 정렬되어 있고 sumB는 내림차순으로 정렬되어 있기 때문에, ptA가 한칸씩 전진하면 그 합이 커질 것이고 ptB가 한칸씩 전진하면 그 합은 작아질 것이다.

따라서, ptA가 가리키는 값인 currentA와 ptB가 가리키는 값인 currentB의 합을 T와 비교하여 처리한다.

그 합이 T와 같은 경우에는 sumA와 sumB에 현재와 같은 수가 있는지 각각 확인하고 count를 증가시킨다.

int ptA = 0;
int ptB = 0;
      
while (ptA < sumA.size() && ptB < sumB.size()) {
    long currentA = sumA.get(ptA);
    long currentB = sumB.get(ptB);

    if (currentA + currentB > T) {
        ptB++;
    }
    else if (currentA + currentB < T) {
        ptA++;
    }
    else {
        long countA = 0;
        long countB = 0;

        // 같은 수가 있는지 확인하고 있다면 그만큼 countA를 증가시킨다
        while (ptA < sumA.size() && sumA.get(ptA) == currentA) {
            ptA++;
            countA++;
        }

        // 같은 수가 있는지 확인하고 있다면 그만큼 countB를 증가시킨다
        while (ptB < sumB.size() && sumB.get(ptB) == currentB) {
            ptB++;
            countB++;
        }

        // A 배열에서 currentA과 같은 수 * B 배열에서 currentB과 같은 수
        answer += countA * countB;
    }
}

 

전체 코드

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

public class Main {

    static long T;
    static int n, m;
    static long[] arrA, arrB;
    static long answer;

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

        StringTokenizer st = new StringTokenizer(br.readLine());
        T = Long.parseLong(st.nextToken());

        st = new StringTokenizer(br.readLine());
        n = Integer.parseInt(st.nextToken());

        st = new StringTokenizer(br.readLine());
        arrA = new long[n];
        for (int i = 0; i < n; i++) {
            arrA[i] = Long.parseLong(st.nextToken());
        }

        st = new StringTokenizer(br.readLine());
        m = Integer.parseInt(st.nextToken());

        st = new StringTokenizer(br.readLine());
        arrB = new long[m];
        for (int i = 0; i < m; i++) {
            arrB[i] = Integer.parseInt(st.nextToken());
        }

        List<Long> sumA = new ArrayList<>();
        List<Long> sumB = new ArrayList<>();

        for (int i = 0; i < n; i++) {
            long sum = arrA[i];
            sumA.add(sum);
            for (int j = i + 1; j < n; j++) {
                sum += arrA[j];
                sumA.add(sum);
            }
        }

        for (int i = 0; i < m; i++) {
            long sum = arrB[i];
            sumB.add(sum);
            for (int j = i + 1; j < m; j++) {
                sum += arrB[j];
                sumB.add(sum);
            }
        }

        Collections.sort(sumA);
        Collections.sort(sumB, Comparator.reverseOrder());

        int ptA = 0;
        int ptB = 0;

        while (ptA < sumA.size() && ptB < sumB.size()) {
            long currentA = sumA.get(ptA);
            long currentB = sumB.get(ptB);

            if (currentA + currentB > T) {
                ptB++;
            }
            else if (currentA + currentB < T) {
                ptA++;
            }
            else {
                long countA = 0;
                long countB = 0;

                // 같은 수가 있는지 확인하고 있다면 그만큼 countA를 증가시킨다
                while (ptA < sumA.size() && sumA.get(ptA) == currentA) {
                    ptA++;
                    countA++;
                }

                // 같은 수가 있는지 확인하고 있다면 그만큼 countB를 증가시킨다
                while (ptB < sumB.size() && sumB.get(ptB) == currentB) {
                    ptB++;
                    countB++;
                }

                // A 배열에서 currentA과 같은 수 * B 배열에서 currentB과 같은 수
                answer += countA * countB;
            }
        }

        System.out.println(answer);
    }
}