[C++] 백준 2143: 두 배열의 합

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

 

시간: 1시간 11분

해결: X

---

문제

두 정수 배열 A, B의 연속 부분합을 더해서 T와 같은 경우의 수를 구해야 합니다.

저는 문제를 두 개로 나눴습니다.

 

모든 부분합 구하기

A의 모든 부분합 SA의 크기는 1+2+3+ ... +n이므로 아래와 같습니다.

B의 모든 부분합 SB도 똑같습니다.

 

부분합 더하기

SA와 SB의 원소를 직접 더 해봅니다.

다만 모든 경우의 수를 전부 구하면 대충 O(n^2*m^2)에 가까워지니 방법이 필요합니다.

 

제가 시도한 방법은

1. 모든 SA의 원소 a에 대해 합이 T가 되는 SB의 원소 b를  이분 탐색

2. lower_bound()로 a+b=T인 가장 작은 인덱스의 원소를 구하고 하나씩 세기

3. (채택) equal_range()로 SB에서 a+b=T인 b의 개수 구하기

 

https://velog.io/@cowkjw/C-STL-upperbound-lowerbound-equalrange-binarysearch

[C++ STL] upper_bound, lower_bound, equal_range, binary_search

equal_range()에 넘긴 범위는 오름차순으로 정렬 되어있어야 합니다.

---

코드

const int SZ = 1001;

int n, m;
int A[SZ], B[SZ];
long long T;
vector<long long> SA, SB;

long long solve(){
    // vector 메모리 확보
    SA.reserve(1LL*n*(n+1)/2);
    SB.reserve(1LL*m*(m+1)/2);
    
    // 모든 부분합 구하기
    for(int i=0; i<n; i++){
        long long sum=0;
        for(int j=i; j<n; j++){
            sum += A[j];
            SA.push_back(sum);
        }
    }
    for(int i=0; i<m; i++){
        long long sum=0;
        for(int j=i; j<m; j++){
            sum += B[j];
            SB.push_back(sum);
        }
    }
    
    // 이분 탐색 위해 정렬
    sort(SA.begin(), SA.end());
    sort(SB.begin(), SB.end());
    
    long long ret=0;
    for(long long a : SA){
        long long target = T-a;
        auto range = equal_range(SB.begin(), SB.end(), target);
        ret += (range.second - range.first);
    }
    
    return ret;
}

reserve()는 vector가 사용할 메모리를 미리 할당해줍니다.

 

vector는 push_back()으로 할당된 메모리보다 크기가 커질 경우, 현재 메모리의 해제하고 그 두 배의 공간을 새로 할당합니다.

이 작업은 최대 O(n)의 시간을 소모하고, 사용할 메모리의 크기를 안다면 reserve()로 사전에 처리하여 시간을 절약할 수 있습니다.

문제에서는 대략 1,000,000개의 원소 공간을 사용하겠네요.