[백준(Python/Java)] 2143_두 배열의 합

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

2143번: 두 배열의 합

 

이 문제에서 구하고자 하는 내용은 아래와 같습니다.

 

한 배열 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, B의 원소의 개수는 1,000개 이하이라고 주어집니다.

 

그렇다면 A의 부분합과 B의 부분합 리스트를 만드는데 우선 N^2 + M^2의 시간이 소요되겠죠?  이 상황에서 각 리스트의 모든 원소들을 비교하면서 합 T를 만들 수 있는가를 확인하려면 N^2*N^2 = N^4의 시간이 소요됩니다. 따라서 문제의 제한시간을 초과할 수 밖에 없는 시간복잡도를 가집니다.

 

따라서, 이 문제는 이분탐색을 활용한 완전탐색을 해주어야 합니다. 즉, 먼저 구하고자 하는 합 T에서 A 부분합 리스트의 각 원소(x)를 제외해보고, 이 때 남은 값(T-x)이 B 부분합 리스트에 몇 개 존재하는지를 세주면 됩니다.

 

파이썬에서는 bisect 라이브러리가 존재하기 때문에 배열에서 Target값이 몇 개 존재하는지를 쉽게 구할 수 있습니다. bisect_left는 정렬된 배열에서 target이 들어갈 수 있는 가장 작은 인덱스를 리턴하고, bisect_right는 가장 큰 인덱스를 리턴하기 때문이죠.

 

예를 들어, 한 배열이 [1, 4, 5, 8, 8, 9]이고 target = 8이라면, bisect_left(array, target) = 3, bisect_right(array, target) = 5가 됩니다. 따라서 두 인덱스의 차이를 구해주면 해당 배열에서 target값이 몇 개 존재하는지 알 수 있죠.

 

다만, 자바에는 bisect 라이브러리가 존재하지 않기 때문에 이분탐색 알고리즘을 직접 구현하여 위와 같은 로직을 수행해주면 됩니다.

 

* answer는 가능한 경우의 수가 백만 * 백만이기 때문에 long으로 선언해줘야 합니다.

 

풀이는 아래와 같습니다.

[풀이 - 파이썬]

from bisect import bisect_left, bisect_right

def getParitalSumArray(array):
    result = []
    length = len(array)
    for i in range(length):
        for j in range(i, length):
            temp = sum(array[i:j+1])
            result.append(temp)
    return sorted(result)

def binarySearch(array, target):
    # target이 위치할 수 있는 왼쪽과 오른쪽 인덱스를 구한다.
    left_index = bisect_left(array, target)
    right_index = bisect_right(array, target)
    count = right_index - left_index
    return count

t = int(input())
n = int(input())
array_A = list(map(int, input().split()))
m = int(input())
array_B = list(map(int, input().split()))

answer = 0
sumOfArray_A = getParitalSumArray(array_A)
sumOfArray_B = getParitalSumArray(array_B)

for x in sumOfArray_A:
    # 전체 합에서 A배열의 합을 뺀 값이 B배열에 몇 개 존재하는가?
    rem = t - x
    answer += binarySearch(sumOfArray_B, rem)

print(answer)

 

[풀이 - 자바]

package PS;

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.lang.reflect.Array;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.StringTokenizer;

public class B_2143 {
    static int t,n,m;
    static int[] array_A, array_B;

    public static ArrayList<Integer> getPartialSumArray(int[] array){
        int length = array.length;
        ArrayList<Integer> result = new ArrayList<>();

        for(int i = 0; i < length; i++){
            int temp = array[i];
            result.add(temp);
            for(int j = i+1; j < length; j++){
                temp += array[j];
                result.add(temp);
            }
            temp = 0;
        }
        Collections.sort(result);
        return result;
    }

    public static int lowerBound(ArrayList<Integer> list, int target){
        int left = 0, right = list.size();
        while(left < right){
            int mid = (left + right) / 2;
            if(list.get(mid) >= target){
                right = mid;
            }else{
                left = mid + 1;
            }
        }
        return right;
    }

    public static int upperBound(ArrayList<Integer> list, int target){
        int left = 0, right = list.size();
        while(left < right){
            int mid = (left + right) / 2;
            if(list.get(mid) <= target){
                left = mid + 1;
            }else{
                right = mid;
            }
        }
        return right;
    }

    public static void main(String[] args) throws Exception {
        // 인풋 입력받기
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        t = Integer.parseInt(br.readLine());
        n = Integer.parseInt(br.readLine());

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

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

        // 부분합 배열을 구하고 오름차순 정렬
        ArrayList<Integer> a = getPartialSumArray(array_A);
        ArrayList<Integer> b = getPartialSumArray(array_B);


        long answer = 0;
        for(int x : a){
            // 전체 합에서 A 부분합을 뺀 값이 B 리스트에 몇 개 존재하는가?
            int rem = t - x;
            answer += upperBound(b, rem) - lowerBound(b, rem);
        }
        System.out.println(answer);
    }
}