문제

2143번: 두 배열의 합

풀이

모든 경우의 부 배열의 합을 A B 각각 구하고 두 부배열의 합이 T인 경우를 찾아주면 된다.
합이 T인 경우를 찾기 위해 이분탐색을 이용해도 되고, 투 포인터를 이용해도 된다.

투 포인터

A, B 부배열의 합을 구해둔 배열을 오름차순으로 정렬한다.
A 부배열의 합 배열은 왼쪽부터(0), B 부배열의 합 배열은 오른쪽부터(.size()-1) 시작하여 각 인덱스에 해당하는 두 수의 합을 구한 뒤 t와 비교한다.
합이 t보다 작다면 A 배열의 인덱스를 +1, 크다면 B 배열의 인덱스를 -1 해주고 다시 비교한다.
합이 t와 같다면 같은 수의 개수를 센 뒤 두 배열에서 구한 같은 수의 개수를 서로 곱해준 뒤 최종 cnt에 더해준다.
인덱스가 배열 범위를 벗어날 때까지 반복한다.

이분 탐색

두 배열의 합 배열을 정렬한 뒤 A 배열에 모든 원소에 대해서 B + t 의 값과 비교하는 이분탐색으로 Upper bound와 Lower bound의 차이를 구해 cnt에 더해준다.

메모리가 64MB로 제한되어 있어서 바짝 쫄았다. 하지만 그럴 필요가 없었다.

코드

투 포인터

#include <iostream>
#include <vector>
#include <algorithm>

using namespace std;

int main()
{
    int t, n ,m;
    vector<int64_t> sumA;
    vector<int64_t> sumB;
    vector<int> a {0};
    vector<int> b {0};

    // 입력
    cin >> t;
    cin >> n;
    for (int i=0; i<n; ++i)
    {
        int temp;
        cin >> temp;
        a.push_back(temp + a[i]);
    }
    cin >> m;
    for (int i=0; i<m; ++i)
    {
        int temp;
        cin >> temp;
        b.push_back(temp + b[i]);
    }

    // 누적합 계산
    for (int i=0; i<n+1; ++i)
        for (int j=0; j<i; ++j)
            sumA.push_back(a[i] - a[j]);

    for (int i=0; i<m+1; ++i)
        for (int j=0; j<i; ++j)
            sumB.push_back(b[i] - b[j]);
    
    // 정렬
    sort(sumA.begin(), sumA.end());
    sort(sumB.begin(), sumB.end());

    // 투포인터 탐색
    int left = 0;
    int right = sumB.size()-1;
    int64_t cnt = 0;
    while (0 <= right && left < sumA.size())
    {
        if (sumA[left] + sumB[right] > t)
            --right;
        
        else if (sumA[left] + sumB[right] < t)
            ++left;
        
        else
        {
            int64_t tempA = sumA[left];
            int64_t tempB = sumB[right];
            int64_t cntA = 0;
            int64_t cntB = 0;

            while (tempA == sumA[left] && left < sumA.size())
            {
                ++left;
                ++cntA; 
            }

            while (tempB == sumB[right] && 0 <= right)
            {
                --right;
                ++cntB;
            }
            cnt += cntA * cntB;
        }
    }

    cout << cnt << endl;
    return 0;
}

이분 탐색

#include <iostream>
#include <vector>
#include <algorithm>

using namespace std;

int main()
{
    int t, n ,m;
    vector<int64_t> sumA;
    vector<int64_t> sumB;
    vector<int> a {0};
    vector<int> b {0};

    // 입력
    cin >> t;
    cin >> n;
    for (int i=0; i<n; ++i)
    {
        int temp;
        cin >> temp;
        a.push_back(temp + a[i]);
    }
    cin >> m;
    for (int i=0; i<m; ++i)
    {
        int temp;
        cin >> temp;
        b.push_back(temp + b[i]);
    }

    // 누적합 계산
    for (int i=0; i<n+1; ++i)
        for (int j=0; j<i; ++j)
            sumA.push_back(a[i] - a[j]);

    for (int i=0; i<m+1; ++i)
        for (int j=0; j<i; ++j)
            sumB.push_back(b[i] - b[j]);
    
    // 정렬
    sort(sumA.begin(), sumA.end());
    sort(sumB.begin(), sumB.end());

    
    int64_t cnt = 0;
    for (auto A : sumA)
    {
        int lowerBound;
        {
            int lo = 0;
            int hi = sumB.size()-1;
            while (lo <= hi)
            {
                int mid = (lo + hi) / 2;

                if (t - A <= sumB[mid])
                    hi = mid - 1;

                else
                    lo = mid + 1;
            }
            lowerBound = lo;
        }

        int upperBound;
        {
            int lo = 0;
            int hi = sumB.size()-1;
            while (lo <= hi)
            {
                int mid = (lo + hi) / 2;

                if (t - A < sumB[mid])
                    hi = mid - 1;

                else
                    lo = mid + 1;
            }
            upperBound = lo;
        }

        cnt += upperBound - lowerBound;
    }

    cout << cnt << endl;
    return 0;
}