백준 / 7453 합이 0인 네 정수 C++

www.acmicpc.net/problem/7453

7453번: 합이 0인 네 정수

입력으로 주어진 A, B, C, D 배열에 대해 4개 수의 합이 0이 되는 순서쌍의 개수를 구하는 문제입니다.

 

처음 생각한 것

1. A, B의 가능한 합을 모두 구한다. ( AB 배열에 저장 )

2. C, D의 가능한 합을 모두 구한다. ( CD 배열에 저장 )

3. C, D를 정렬한다.

4. binary search를 이용해 AB의 모든 원소를 key로 하여 CD 배열에서 key의 index를 찾는다.

5. 찾은 경우 해당 index의 앞 뒤로 같은 것의 개수를 카운트한다.

=> 시간 초과

 

시간 초과를 해결하기 위해 lower_bound와 upper_bound를 이용했습니다.

(직접 구현하니 70% 쯤에서 시간초과가 뜬다....)

=> key를 기준으로 앞뒤로 같은 것 찾는 과정을 안 해도 됨. algorithm의 lower_bound, upper_bound 이용하여 key에 근접한 index를 구한다.

 

lower_bound : key와 같은 값 있으면 가장 작은 index를 구한다. 없으면 key보다 큰 수 중 가장 작은 원소의 index를 구한다.

 

upper_bound : key보다 큰 원소 중 가장 작은 원소의 index를 구한다.

 

코드는 다음과 같습니다.

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

using namespace std;

// A,B 와 C,D 의 합을 계산한다.
// C,D의 합만 오름차순으로 정렬한다.
// lower_bound, upper_bound 함수를 이용해 -num 에 가까운 index를 찾는다.
// lower_bound : key와 같은 값 있으면 가장 작은 index를 구한다.
// 없으면 key보다 큰 수 중 가장 작은 원소의 index를 구한다.
// upper_bound : key보다 큰 원소 중 가장 작은 원소의 index를 구한다.
// key (여기서는 -num) 와 같은 값 존재할 때 조건 만족하므로 upper - lower를 더한다.

int main()
{
	ios::sync_with_stdio(false);
	cin.tie(NULL);
	cout.tie(NULL);

	int n;
	long long arr[4000][4];
	long long result = 0;
	vector<long long>cd;

	cin >> n;

	for (int i = 0; i < n; i++) {
		for (int j = 0; j < 4; j++) {
			cin >> arr[i][j];
		}
	}

	for (int i = 0; i < n; i++) {
		for (int j = 0; j < n; j++) {
			cd.push_back(arr[i][2] + arr[j][3]);
		}
	}

	sort(cd.begin(), cd.end());

	for (int i = 0; i < n; i++) {
		for (int j = 0; j < n; j++) {
			long long num = arr[i][0] + arr[j][1];
			long long lower = lower_bound(cd.begin(), cd.end(), (-1) * num) - cd.begin();
			long long upper = upper_bound(cd.begin(), cd.end(), (-1) * num) - cd.begin();

			if ((-1) * num == cd[lower])
				result += (upper - lower);
		}
	}

	cout << result;
}