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

문제

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

---

문제 요약

수열 A, B가 주어질 때, A의 부 배열과 B의 부 배열 합이 T가 되도록하는 모든 부 배열 쌍의 수를 출력한다.

 

A의 부 배열은 A[i] + ... + A[j], B의 부 배열은 B[i] + ... + B[j]를 의미한다.

 

n은 1이상 1000이하이다. 

---

코드

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

using namespace std;

int t, n, m;

int A[1010];
int B[1010];

int prefix_A[1010];
int prefix_B[1010];

vector<int> sub_A, sub_B;

int main() {
    // 배열 A, B 입력
    cin >> t >> n;

    for (int i = 1; i <= n; i++) {
        cin >> A[i];
        prefix_A[i] = prefix_A[i - 1] + A[i];
    }

    cin >> m;

    for (int i = 1; i <= m; i++) {
        cin >> B[i];
        prefix_B[i] = prefix_B[i - 1] + B[i];
    }

    // 배열의 각 부분 합 계산
    for (int i = 0; i <= n; i++) {
        for (int j = i + 1; j <= n; j++) {
            sub_A.push_back(prefix_A[j] - prefix_A[i]);
        }
    }

    for (int i = 0; i <= m; i++) {
        for (int j = i + 1; j <= m; j++) {
            sub_B.push_back(prefix_B[j] - prefix_B[i]);
        }
    }

    // B의 부분합 배열 정렬
    sort(sub_B.begin(), sub_B.end());

    long long res = 0;

    for (int i = 0; i < sub_A.size(); i++) {
        // 합이 t가 되기 위해 sub_B에서 필요한 수
        int target = t - sub_A[i];

        int lo = lower_bound(sub_B.begin(), sub_B.end(), target) - sub_B.begin();
        int hi = upper_bound(sub_B.begin(), sub_B.end(), target) - sub_B.begin();

        res += (hi - lo);
    }

    cout << res;


    return 0;
}

---

코드 설명

1. 배열 A, B의 값을 입력 받고, 빠르게 부 배열의 합을 구하기 위해 누적합 배열을 미리 구한다.

 

2. 누적합을 이용해 A와 B에서 각각 가능한 부 배열의 합을 구해 벡터 sub_A, sub_B에 저장한다.

(벡터의 크기는 n × (n + 1) / 2이 된다.)

 

3. lower_bound와 upper_bound 사용을 위해, 부 배열의 합 배열 하나를 정렬한다.

(sub_B를 정렬한다.)

 

4. 경우의 수를 모두 탐색하기 위해, 먼저 sub_A에서 원소 하나를 고르고, T에서 뺀다.

(우리는 sub_B에서 T - sub_A를 찾으면 된다.)

 

5. lower_bound와 upper_bound를 사용하여 sub_B에서 T - sub_A의 개수를 이진 탐색을 통해 구한다.

 

6. 모든 경우의 수에서 구한 T - sub_A의 개수의을 출력한다.

---

고찰

lower_bound(배열 시작 주소, 배열 끝 주소, target)

=> 배열 범위의 원소들 중 target 값보다 크거나 같은 첫 번째 원소의 주소 반환

 

upper_bound(배열 시작 주소, 배열 끝 주소, target)

=> 배열 범위의 원소들 중 target 값보다 큰 첫 번째 원소의 주소 반환

 

반환 값에서 배열의 시작 주소를 빼주면, 해당 인덱스 번호를 알 수 있다.

 

ex)

lower_bound(sub_B.begin(), sub_B.end(), target) - sub_B.begin();

 

따라서, 하나의 값이 배열에 1개 이상 존재할 때, 시작 인덱스와 끝 인덱스를 알 수 있다.

 

추가로 이를 통해 해당 값이 배열에 몇 개 존재하는 지도 알 수 있다.

 

시간 복잡도는 배열을 정렬하는데 O(n log n),  함수를 호출할 때마다 O(log n)이 걸린다.

(이진 탐색을 이용하기 때문이다.)