[백준 2143 java 자바]두 배열의 합 - 투포인터(부분합, 이중배열)

[Gold III] 두 배열의 합 - 2143

문제 링크

성능 요약

메모리: 50704 KB, 시간: 920 ms

분류

이분 탐색(binary_search), 누적 합(prefix_sum)

문제 설명

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

---

풀이

부배열에 대한 개념을 파악하는게 먼저다.

부배열은 배열에서 연속된 수의 부분합이다.(ex, A[1] + A[3]은 안되고, A[1]+A[2] 가능)

두 배열의 각각 부배열 합의 List를 구한다음, 투 포인터를 돌려서 카운팅을 하면 된다.

투포인터는 다음 블로그를 참고하면 좋을 것 같다. https://lotuslee.tistory.com/30?category=963570

코드

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.*;
import java.util.stream.Collectors;
import java.util.stream.IntStream;
import java.util.stream.Stream;

import static java.util.Arrays.stream;

class Main {
    static int[] A,B;

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        int T = Integer.parseInt(br.readLine());
        int n = Integer.parseInt(br.readLine());
        A = stream(br.readLine().split(" "))
                .mapToInt(Integer::parseInt).toArray();
        int m = Integer.parseInt(br.readLine());
        B = stream(br.readLine().split(" "))
                .mapToInt(Integer::parseInt).toArray();

        //초기화


        //부배열 쌍
        ArrayList<Integer> accumA = new ArrayList<>();
        ArrayList<Integer> accumB = new ArrayList<>();


        //부배열 구하기
        for (int i = 0; i < n; i++) {
            int sum = 0;
            for(int j = i; j < n; j++){
                sum += A[j];
                accumA.add(sum);
            }
        }
        for (int i = 0; i < m; i++) {
            int sum = 0;
            for (int j = i; j < m; j++) {
                sum += B[j];
                accumB.add(sum);
            }
        }

        //부배열 정렬
        Collections.sort(accumA);
        Collections.sort(accumB);


        //이중배열 투포인터  
        int aSize = accumA.size();
        int bSize = accumB.size();

        //왼쪽은 왼쪽부터, 오른쪽은 오른쪽부터 배열을 진행한다.
        long cnt = 0;
        int leftPointer = 0;
        int rightPointer = accumB.size() - 1;

        //정렬을 했기 때문에, 왼쪽 + 오른쪽 합이 크면 오른쪽에서 <- , 작으면 왼쪽에서 -> 으로 가서 크기를 비교한다.
        while (leftPointer < aSize && rightPointer >= 0) {
            int sum = accumA.get(leftPointer) + accumB.get(rightPointer);

            //같을시에 정렬을 했기 때문에 같은 수는 일렬로 쭈르륵 있을것이다.
            //그럼 연속된 같은 수는 묶어서 A배열의 개수 * B 배열의 개수 를 더해 counting 해준다.
            //ex A[1,1,1,2,3,4], B[1,3,3,4] 일 경우, 합 4를 찾으면
            //ACnt = 연속된 1이 3개 = 3, BCnt = 연속된 3이 2개 = 2
            //Cnt += 3*2 (6가지 경우가 나오므로)
            if (sum == T) {

                int a = accumA.get(leftPointer);
                int b = accumB.get(rightPointer);
                long leftCnt = 0;
                long rightCnt = 0;

                while(leftPointer < aSize && accumA.get(leftPointer) == a){
                    leftCnt++;
                    leftPointer++;
                }
                while(rightPointer >= 0 && accumB.get(rightPointer) == b){
                    rightCnt++;
                    rightPointer--;
                }
                cnt += leftCnt * rightCnt;
                //작으면 -> 가서 크기를 키운다
            } else if (sum < T) {
                leftPointer++;
                //크면 <- 가서 크기를 줄인다.
            } else if (sum > T) {
                rightPointer--;
            }
        }


        System.out.println(cnt);


    }

}