[백준] 두 배열의 합

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

2143번: 두 배열의 합

우선 문제부터 간단하게 요약하면, 

 

A={1, 3, 1, 2}, B={1, 3, 2}, T = 5 이렇게 A, B 배열과 T값이 주어질 때, A의 부배열과 B의 부배열의 합이 T가 되도록 하는 모든 경우의 수의 갯수를 구해야 합니다. 부 배열과 같은 부분에 대한 정의는 문제를 한번 직접 읽어보는 것을 추천합니다.

 

<Solution>

사실 쓰이는 알고리즘 자체는 크게 어렵지 않습니다. 

1. A 배열과 B 배열의 Prefix sum array를 계산합니다. 2. Prefix sum array를 사용하여, 부배열으로 가능한 값들의 합을 모두 구합니다. 3. B의 부배열으로 가능한 값 array를 정렬합니다. (필자의 코드에서는 candidateB) 4. candidateA array의 각 원소에 대해 T와의 차이값을 binary lower upper bound를 이용하여 candidateB에서 찾습니다.

위와 같은 과정으로 문제를 해결할 수 있습니다. 

 

기본적으로 Prefix sum array를 해야하는 이유는 우리는 부배열으로 가능한 모든값들을 미리 계산해 둘 것인데, Prefix sum을 이용하지 않고 매번 더해나가는 과정을 진행하면 같은 계산을 여러번 진행 하게 되어 시간이 오래 걸릴 것이기 때문입니다.

 

추가적으로 binary searching은 upper bound와 lower bound를 찾을 때 사용할 수 있다는 점을 활용하면 candidateB array에서 원하는 값을 빠르게 찾아낼 수 있기에 문제를 빠른 시간내에 해결할 수 있습니다. 

 

실제 구현은 아래와 같습니다.

#include <iostream>
#include <algorithm>

using namespace std;

int T;
int n, m;
int A[1000+10];
int B[1000+10];

int prefixA[1000+10];
int prefixB[1000+10];

int candidateA[1000*1000];
int candidateB[1000*1000];

// find value in candidateB
int binaryLowerBound(int toFind, int s, int e){
    int sol = -1;

    while(e>=s){
        int mid = (s+e) / 2;

        if(candidateB[mid] == toFind){
            e = mid-1;
            sol = mid;
        }
        else if(candidateB[mid] > toFind){
            e = mid-1;
        }
        else{
            s = mid+1;
        }
    }
    return sol;
}

int binaryUpperBound(int toFind, int s, int e){
    int sol = -1;
    
    while(e>=s){
        int mid = (s+e) / 2;

        if(candidateB[mid] == toFind){
            s = mid+1;
            sol = mid;
        }
        else if(candidateB[mid] > toFind){
            e = mid-1;
        }
        else{
            s = mid+1;
        }
    }
    return sol;
}

int main(){

    // get input
    cin >> T >> n;
    for (int i = 1; i<=n; i++){
        cin >> A[i];
    }
    cin >> m;
    for(int i = 1; i<=m; i++){
        cin >> B[i];
    }

    // build prefix sum array
    for(int i = 1; i<=n; i++){
        prefixA[i] = prefixA[i-1] + A[i];
    }
    for(int i = 1; i<=m; i++){
        prefixB[i] = prefixB[i-1] + B[i];
    }

    // find all candidates and save it on array
    int k1 = 0;
    for(int i = 0; i<=n-1; i++){
        for(int j = i+1; j<=n; j++){
            candidateA[k1] = prefixA[j] - prefixA[i];
            k1++;
        }
    }
    int k2 = 0;
    for(int i = 0; i<=m-1; i++){
        for(int j = i+1; j<=m; j++){
            candidateB[k2] = prefixB[j] - prefixB[i];
            k2++;
        }
    }

    // sort B for binary search
    sort(candidateB, candidateB+k2);

    // search values for each case
    long long ans = 0;

    for(int i = 0; i<k1; i++){
        int toFind = T-candidateA[i];

        // binary search lower bound && upper bound 
        int lowerBound = binaryLowerBound(toFind, 0, k2-1);
        if(lowerBound == -1) continue; // can't find
        int upperBound = binaryUpperBound(toFind, lowerBound, k2-1);

        ans += upperBound - lowerBound +1;
    }

    cout << ans << endl;
}