백준 2143: 두 배열의 합[Java] - 포포

문제

한 배열 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을 출력한다.

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

 

---

1. 먼저 배열 두 개를 각각 입력받는다.
2. 반복문을 사용해서 부배열을 생성한다(ArrayList).
3. 투포인터를 사용해 부배열의 합과 목표 합인 T를 비교한다.

문제의 핵심은 투포인터를 사용해서 부배열의 합을 구하는 과정이다. 
예시처럼  A = {1, 3, 1, 2}, B = {1, 3, 2}인 경우
A의 부배열은 subA ={1, 3, 1, 2, 4, 4, 3, 5, 6, 7}
B의 부배열은 subB ={1, 3, 2, 4, 5, 6}으로 구할 수 있다.

투포인터 사용을 위해 subA는 오름차순, subB는 내림차순으로 정렬하면
subA ={1, 1, 2, 3, 3, 4, 4, 5, 6, 7}
subB= {6, 5, 4, 3, 2, 1} 로 정렬된다.

이제 ptA, ptB = 0 부터 시작해 두 부배열의 합을 구한다 

subA.get(ptA) + subB.get(ptB)

그리고 경우의 수는 세가지로 나뉜다.

1. 두 부배열의 합이 T보다 큰경우
2. 부배열의 합이 T보다 작은경우
3. 부배열의 합이 T인경우

만약 1번 경우처럼 부배열의 합이 T보다 크다면 ?  ptB를 1 증가시키면 된다. subB는 내림차순으로 정렬했기 때문
반대로 부배열의 합이 T보다 작다면 오름차순 정렬한 subA의 커서를 1 증가시킨다.

만약 부배열의 합이 T라면 현재 ptA와 ptB가 가리키는 부배열과 같은 원소가 잇달아 몇 번 등장하는지 subA와 subB를 순서대로 뒤지면서 개수를 카운트하면 된다.

제출한 코드에서는 두 부배열의 합 == T를 사용하기보다는
target이라는 변수를 사용했다.  

target = T - subA.get(ptA); 
결국 subB.get(ptB)를 target으로 치환했다고 생각해도 좋다. 
코드를 보면 이해할 수 있을 것이다.

제출한 코드

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

public class Main {

    static long T;
    static int N, M;
    static long[] inputA, inputB;

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

        T = Long.parseLong(br.readLine());

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

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

        List<Long> subA = new ArrayList<>();
        List<Long> subB = new ArrayList<>();

        //subA 구하기
        for (int i = 0; i < N; i++) {
            long sum = inputA[i];
            subA.add(sum);
            for (int j = i + 1; j < N; j++) {
                sum += inputA[j];
                subA.add(sum);
            }
        }
        //subB 구하기
        for (int i = 0; i < M; i++) {
            long sum = 0;
            for (int j = i; j < M; j++) {
                sum += inputB[j];
                subB.add(sum);
            }
        }

        //subA, subB 정렬하기
        Collections.sort(subA);
        Collections.sort(subB, Comparator.reverseOrder());

        long result = 0;
        int ptA = 0;
        int ptB = 0;
        while (true) {
            long currentA = subA.get(ptA);
            long target = T - currentA;
            //currentB == target -> subA, subB 같은 수 개수 체크 -> 답구하기. ptA+ ptB+
            if (subB.get(ptB) == target) {
                long countA = 0;
                long countB = 0;
                while (ptA < subA.size() && subA.get(ptA) == currentA){
                    countA++;
                    ptA++;
                }
                while (ptB < subB.size() && subB.get(ptB) == target){
                    countB++;
                    ptB++;
                }
                result += countA * countB;
            }
            //currentB > target -> ptB++
            else if (subB.get(ptB) > target) {
                ptB++;
            }
            //currentB < target -> ptA++
            else {
                ptA++;
            }

            //탈출 조건
            if (ptA == subA.size() || ptB == subB.size()) {
                break;
            }
        }
    }
}