[BOJ] 백준 2143 두 배열의 합 c++ (이분탐색, 누적 합)

문제 출처 : https://www.acmicpc.net/problem/2143

2143번: 두 배열의 합

풀이

어떻게 이분탐색을 돌려야할지 고민이 되는 문제이다.

중요한 부분은, A와 B의 누적합을 미리 저장해두는 것이다.

 

누적 합

A, B 배열에서 모든 부 배열의 누적합을 저장한다.

vector<int> aSum, bSum;

for (int i = 0; i < N; i++) {
    int sum = A[i];
    aSum.push_back(sum);
    for (int j = i + 1; j < N; j++) {
        sum += A[j];
        aSum.push_back(sum);
    }
}

... b도 반복

 이렇게 되면, aSum, bSum에 각각 부 배열의 누적합들이 저장될 것이다.

lower_bound, upper_bound

부 배열의 쌍을 구하는 것이 정답이다. 즉 aSum의 원소와 bSum의 원소를 더했을때 T가 나오면 된다.

따라서 T - aSum의 원소 = bSum 원소 이므로, T - aSum에 해당하는 bSum의 개수를 구해주면 된다.

 

그렇다면 bSum의 개수는 어떻게 구할까?

바로 lower_bound, upper_bound를 활용한다.

lower_bound는 배열에서 범위 내의 원소들 중 value값 보다 크거나 같은 첫 번째 원소의 주소를 리턴

upper_bound는 배열에서 처음으로 value값을 초과하는 원소의 주소를 리턴

즉, 해당 숫자의 시작과 끝 주소를 알 수 있다.

sort(bSum.begin(), bSum.end());
long long ans = 0;
for (int i = 0; i < aSum.size(); i++) {
    int target = T - aSum[i];
    // 타겟과 같은 수들 개수 (hi - lo)
    int lo = lower_bound(bSum.begin(), bSum.end(), target) - bSum.begin();
    int hi = upper_bound(bSum.begin(), bSum.end(), target) - bSum.begin();
    ans += (hi - lo);
}

lo부터 hi까지 해당 수이므로 hi - lo를 더해준다.

---

코드

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

using namespace std;

int T, N, M;
int A[1001], B[1001];

int main(void){
    ios::sync_with_stdio(0); cin.tie(0); cout.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];

	vector<int> aSum, bSum;
	
	for (int i = 0; i < N; i++) {
		int sum = A[i];
		aSum.push_back(sum);
		for (int j = i + 1; j < N; j++) {
			sum += A[j];
			aSum.push_back(sum);
		}
	}

	for (int i = 0; i < M; i++) {
		int sum = B[i];
		bSum.push_back(sum);
		for (int j = i + 1; j < M; j++) {
			sum += B[j];
			bSum.push_back(sum);
		}
	}

	sort(bSum.begin(), bSum.end());
	long long ans = 0;
	for (int i = 0; i < aSum.size(); i++) {
		int target = T - aSum[i];
        // 타겟과 같은 수들 개수 (hi - lo)
		int lo = lower_bound(bSum.begin(), bSum.end(), target) - bSum.begin();
		int hi = upper_bound(bSum.begin(), bSum.end(), target) - bSum.begin();
		ans += (hi - lo);
	}

	cout << ans;

	return 0;
}