백준 10816 숫자 카드 2

www.acmicpc.net/problem/10816

10816번: 숫자 카드 2

문제의 내용은 간단하다. 처음에는 map이나 set을 사용해서 각각의 탐색은 O(log N)이니 찾을 수 있을 것 같았지만

 

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

int main() {
	ios::sync_with_stdio(false); cin.tie(0);
	int n, m, x;
	vector<int> ar;
	map<int, int> M;
	cin >> n;
	for (int i = 0; i < n; i++) {
		cin >> x;
		M[x]++;
	}
	cin >> m;
	for (int i = 0; i < m; i++) {
		cin >> x;
		auto it = M.begin();
		for (it = M.begin(); it != M.end(); it++) {
			if (it->first == x) {
				cout << it->second << ' '; break;
			}
		}
		if (it == M.end()) cout << 0 << ' ';
	}
	return 0;
}

이 코드는 TLE가 났다. 알고 보니 map이나 set은 레드 블랙 트리를 사용하는데 탐색의 시간은 O(log N)이 맞지만 

삽입하는 과정에서 매우 길어지기 때문에 입력받는 곳에서 시간 초과가 난 것이다.

 

이후 풀이법을 찾아봤더니

upper_bound와 lower_bound를 통해 각 인덱스를 알아서 차를 구하면 된다.

 

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

int main() {
	ios::sync_with_stdio(false); cin.tie(0);
	int n, m, x;
	vector<int> ar;
	cin >> n;
	for (int i = 0; i < n; i++) {
		cin >> x; ar.push_back(x);
	}
	sort(ar.begin(), ar.end());
	cin >> m;
	for (int i = 0; i < m; i++) {
		cin >> x;
		int cnt = (upper_bound(ar.begin(), ar.end(), x) - ar.begin()) \
			- (lower_bound(ar.begin(), ar.end(), x) - ar.begin());
		cout << cnt << ' ';
	}
	return 0;
}

정렬된 배열에서 찾으려는 키값의 갯수를 알아내려면 이 방식으로 풀어야겠다.