[백준 2143번 / java] 두 배열의 합

문제

2143번: 두 배열의 합 (acmicpc.net)

1. 배열 A, B에 대해 각 부배열을 구한다.

2. 부배열은 연속된 원소들로 이루어진 부분배열이다.

3. A배열의 부배열의 합 + B배열의 부배열의 합 = T가 되는 모든 부배열 쌍의 개수를 구하라.

풀이

1. 배열의 부분합을 구한다.

2. 부분합을 각 리스트에 넣는다.

3. 두 리스트의 값을 투포인터로 선택한다.

4. 합이 T가 되는 경우를 count한다.

 

부배열의 합을 구할 때 시간복잡도 :

A, B의 원소 개수는 각각 최대 1,000개이다. 

이때 배열의 모든 부배열을 구하려면 부배열의 원소가 1, 2, ..., 1,000개인 경우를 모두 구해야 하므로 시간복잡도는 배열 한 개당 O(n^2)

 

두 리스트에서 값을 하나씩 선택하는 경우 시간복잡도 :

한 리스트에 저장된 모든 부배열의 합은 약 n^2개이다.

브루트포스를 사용하는 경우 -> 총 O(n^2 * n^2) = n^4 최악의 경우 1000,000,000,000 > 2억 (사용 불가)

따라서 리스트를 정렬한 후 투포인터를 사용하는 방법이 적절하다.

투포인터의 경우 시간복잡도는 n이다. -> 총 O(n^2 * n) = n^3 최악의 경우 1000,000,000 < 2억 (사용 가능)

 

코드

import java.util.*;
import java.io.*;
public class Main {
    static long T;
    static int n, m;
    static int[] arr1, arr2;
    static List<Integer> list1, list2;
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
        StringTokenizer st;
        T = Integer.parseInt(br.readLine());

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

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

        list1 = partialSum(arr1);
        list2 = partialSum(arr2);

        Collections.sort(list1);
        Collections.sort(list2);

        bw.write(select(list1, list2) + " ");
        bw.close();

    }
    static List<Integer> partialSum(int[] arr){
        List<Integer> list = new ArrayList<>();
        for(int i=0; i<arr.length; i++){
            int sum = 0;
            for(int j=i; j>=0; j--){
                sum += arr[j];
                list.add(sum);
            }
        }
        return list;
    }
    static long select(List<Integer> list1, List<Integer> list2){
        //이중 포인터로 두 부분집합의 합이 T인 경우를 count
        long count = 0;
        int leftPointer = 0;
        int rightPointer = list2.size() - 1;

        while(leftPointer < list1.size() && rightPointer >= 0){
            long leftNum = list1.get(leftPointer);
            long rightNum = list2.get(rightPointer);
            long sum = leftNum + rightNum;

            if(sum == T){
                long leftSameNumSize = 0;
                long rightSameNumSize = 0;
                while(leftPointer < list1.size() && list1.get(leftPointer) == leftNum){
                    leftSameNumSize++;
                    leftPointer++;
                }
                while(rightPointer >= 0 && list2.get(rightPointer) == rightNum){
                    rightSameNumSize++;
                    rightPointer--;
                }
                count += leftSameNumSize * rightSameNumSize;
            }else if(sum < T){
                leftPointer++;
            }else {
                rightPointer--;
            }
        }
        return count;
    }
}

 

Note

• 처음에 부분합을 구할 때, T보다 작은 합을 가진 부분집합의 합만을  list에 추가해서 틀렸다.

 static List<Integer> partialSum(int[] arr){
        //T보다 작은 합을 가진 부분집합의 합을 모두 list에 추가
        List<Integer> list = new ArrayList<>();
        for(int i=0; i<arr.length; i++){
            int sum = 0;
            for(int j=i; j>=0; j--){
                sum += arr[j];
                if(sum < T){
                    list.add(sum);
                }else break;
            }
        }
        return list;
    }