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

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

2143번: 두 배열의 합

 

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;

int t, n, m;
vector<int> a, b, v, w;

int main() {
	cin >> t;
	cin >> n;
	for(int i = 0; i < n; i++){
		int x;
		cin >> x;
		a.push_back(x);
	}
	cin >> m;
	for(int i = 0; i < m; i++){
		int x;
		cin >> x;
		b.push_back(x);
	}
	
	for(int i = 0; i < n; i++){
		int sum = a[i];
		v.push_back(sum);
		for(int j = i + 1; j < n; j++){
			sum += a[j];
			v.push_back(sum);
		}
	}
	for(int i = 0; i < m; i++){
		int sum = b[i];
		w.push_back(sum);
		for(int j = i + 1; j < m; j++){
			sum += b[j];
			w.push_back(sum);
		}
	}
	
	sort(w.begin(), w.end());
	
	long long ans = 0;
	for(int i = 0; i < v.size(); i++){
		int diff = t - v[i];
		auto ub = upper_bound(w.begin(), w.end(), diff);
		auto lb = lower_bound(w.begin(), w.end(), diff);
		ans += (ub - lb);
	}
	cout << ans;
	return 0;
}

 

 1. vector a, b에 입력을 받는다.

 

 2. a와 b의 누적합을 v와 w에 각각에 저장한다.

 

 3. 이분탐색을 위해 b를 정렬한다.

 

 4. a + b = t이므로 모든 a에 대해 b가 t - a인 값을 찾는다. 해당 값이 여러개 있을 수 있으므로 upper_bound와 lower_bound의 index를 구해 둘의 차이를 ans에 더해준다.

 

* upper_bound : 찾으려는 key 값보다 같거나 큰 숫자가 배열 몇 번째에서 처음 등장하는지 찾음

* lower_bound : 찾으려는 key 값을 초과하는 숫자가 배열 몇 번째에서 처음 등장하는지 찾음