[백준 2143번] 두 배열의 합 (C++)

문제링크 : https://www.acmicpc.net/problem/2143

2143번: 두 배열의 합

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef pair<int, int> pii;
const int MAX = 987654321;

int T, N, M;
int a[1001], b[1001];
vector<int> sum_a, sum_b;
ll result = 0;

int main()
{
    ios_base::sync_with_stdio(0); 
	cin.tie(0);

    cin >> T >> N;
    
    for(int i=0; i<N; i++)
    {
        cin >> a[i];
    }
    cin >> M;
    for(int i=0; i<M; i++)
    {
        cin >> b[i];
        
    }
    
    for(int i=0; i<N; i++)
    {
        int sum = 0;
        for(int j=i; j<N; j++)
        {
            sum+=a[j]; //a 부 배열의 합 저장
            sum_a.push_back(sum);
        }
    }

    for(int i=0; i<M; i++)
    {
        int sum = 0;
        for(int j=i; j<M; j++)
        {
            sum+=b[j]; //b 부 배열의 합 저장
            sum_b.push_back(sum);
        }
    }
    
    sort(sum_a.begin(), sum_a.end()); //이분 탐색을 위해 정렬
    ll result = 0;
    for(int i=0; i<sum_b.size(); i++)
    {
        int diff = T - sum_b[i];
        int low = lower_bound(sum_a.begin(), sum_a.end(), diff)-sum_a.begin(); //diff보다 같거나 큰 숫자 위치 반환
        int high = upper_bound(sum_a.begin(), sum_a.end(), diff)-sum_a.begin(); //diff를 초과하는 숫자 위치 반환
        result += (high-low); //diff 갯수 저장
    }

    cout << result;
    return 0;
}

누적합의 응용 문제이다.

다루는 배열이 2개이므로 그에 대한 누적합 배열도 각각 저장해준다. (sum_a, sum_b)

 

구하고자 하는 것이 각 배열의 누적합 배열의 값을 더해서 T가 되는 것이다. (sum_a[i] + sum_b[i] == T)

이를 이용하여 T - sum_b[i](이하 diff)를 통해서 이에 만족하는 sum_a[i]의 값을 구해주도록 했다.

이분탐색을 사용했으므로 이를 위해 sum_a[]배열을 정렬부터 해준다.

 

lower_bound를 통해서 diff보다 같거나 큰 숫자의 위치를 반환해주고,

upper_bound를 통해서 diff를 초과하는 값의 위치를 반환해준다.

이렇게 구한 high의 위치에서 low의 위치를 빼주면 diff의 갯수가 나온다.

이를 sum_b[]의 크기 만큼 반복해주며 모든 값에 대해 탐색해준다.

 

다른 풀이를 찾아보니 이분탐색이 제일 많긴 했으나 두 포인터, 해시맵등의 풀이를 볼 수 있었다.

나중에 다른 풀이로 다시 풀어보는 것도 좋을 것 같다.