BOJ 2143 두 배열의 합

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

2143번: 두 배열의 합

 

모든 경우의 수를 Naive 하게 구하려고 하면 다음과 같은 시간복잡도가 걸린다.

O (n (n C 2) ^2)

문제를 잘 보면 입력 배열의 "부 배열"을 이용해야 함을 알 수 있다.
따라서 누적합을 미리 구해놓으면 구간 [i,j] 의 구간합을 O(1) 만에 구할 수 있다.
이 경우, 시간복잡도는 다음과 같다.

O ((n C 2) ^2)

다음으로 건너갈 수 있었던 것은 내 경험적인 요소가 크게 작용했다.
예전에 했던 알고리즘 스터디와 과외에서 meet in middle을 배운 적이 있었다. 해당 문제와 매우 유사한 경우에 사용했었기 때문에 이 문제에서도 해당 방법을 사용했다.
meet in middle까지 사용하는 경우, 시간복잡도는 다음과 같다.

O(n log n)

 

 

#include <cstdio>
#include <vector>
#include <algorithm>

using namespace std;

long long A[1000], B[1000];
int la,lb; 
long long SA[1001], SB[1001];
vector<long long> VA, VB;

void set(long long* a, long long* sa, int l, vector<long long>& va){
    for(int i=0;i<l;i++) sa[i+1] = a[i] + sa[i];
    for(int i=1;i<=l;i++){
        for(int j=i;j<=l;j++){
            long long psum = sa[j] - sa[i-1];
            va.push_back(psum);
        }
    }
    
    sort(va.begin(), va.end());
}

int main(){
    long long target; scanf("%lld",&target);
    unsigned long long answer = 0;
    scanf("%d",&la);
    for(int i=0;i<la;i++) scanf("%lld",A+i);
    scanf("%d",&lb);
    for(int i=0;i<lb;i++) scanf("%lld",B+i);

    set(A,SA,la,VA);
    set(B,SB,lb,VB);

    int ia = 0, ib = VB.size()-1;
    while(ia < VA.size() && ib >= 0){
        long long t = VA[ia] + VB[ib];
        if(t == target){
        	unsigned long long cnta = 1, cntb = 1;
        	while(ib >= 1 && VB[ib] == VB[ib-1]) ib--; cntb+=1;
        	while(ia < VA.size() -1 && VA[ia] == VA[ia+1]) ia++; cnta+=1;
        	answer += cnta * cntb;
        	ib--;ia++;
        }
        else{
        	if(t > target)ib--;
        	else ia++;
        }
    }

    printf("%llu\n",answer);
}

 

PS1. 공간복잡도 상수최적화가 가능하다.
PS2. 입력 처리도 set 함수 안에 넣으면 코드 중복을 없앨 수 있다.