[백준BOJ] 3078 좋은 친구

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

3078번: 좋은 친구

---

문제해설

입력의 갯수가 30만이 이기 때문에, 전수조사를 할경우 O(N^2)롤 TLE에 걸리게 됩니다.

O(NlgN) 정도의 알고리즘이 필요해 보입니다.

 

<입력처리>

	// 이름의 길이에 등수를 push_back
	for (int i = 0; i < N; i++) {
		string input; cin >> input;
		vec[input.size()].push_back(i + 1);
	}

입력이 문자열로 들어오지만, 저희가 필요한 정보는 문자열의 길이와 등수뿐 입니다.

 

 

<정렬>

이문제의 핵심이라고 할 수 있는 부분입니다.

이름 길이가 같은 친구끼리 묶인 vector에서 정렬을 통해서

이후에, 선형 조사를 하게 되던, lower_bound, upper_bound를 하게되던 정렬을 해야만 

큰 입력에 대해서 시간안에 답을 구할 수 있습니다.

	// i 길이의 이름을 가진 친구들을 등수별로 sort
	for (int i = 0; i < 21; i++) {
		sort(vec[i].begin(), vec[i].end());
	}

<upper_bound>

i 길이의 이름을 가진 친구중에서 j번째 등수기준으로

j + K 번째 등수까지 몇명이 있는지 알기위해서

이분탐색을통해서 알 수 있습니다.

	for (int i = 0; i < 21; i++) {
		// upper_bound를 편의를 위한 push_back
		vec[i].push_back(1e9);
		for (int j = 0; j < vec[i].size()-1; j++) {
			// 자신의 등수보다 K 보다 큰 마지막 친구가 iter로 반환 됩니다.
			auto iter = upper_bound(vec[i].begin(), vec[i].end(), vec[i][j] + K); iter--;
			// 시작 주소를 빼줌으로써 index값을 얻을 수 있고
			auto index = iter - vec[i].begin();
			// 마지막 친구의 index - 현재 index를 해주면 친구의 명수가 나오게 됩니다.
			ans += index - j;
		}
	}

---

전체코드

#include <iostream>
#include <vector>
#include <stack>
#include <queue>
#include <algorithm>
#include <cmath>
#include <string>
using namespace std;

vector<int> vec[21];

int main() {
	ios::sync_with_stdio(false); cin.tie(0); cout.tie(0);
	int N, K; cin >> N >> K;

	// 이름의 길이에 등수를 push_back
	for (int i = 0; i < N; i++) {
		string input; cin >> input;
		vec[input.size()].push_back(i + 1);
	}

	// i 길이의 이름을 가진 친구들을 등수별로 sort
	for (int i = 0; i < 21; i++) {
		sort(vec[i].begin(), vec[i].end());
	}

	long long ans = 0;
	for (int i = 0; i < 21; i++) {
		// upper_bound를 편의를 위한 push_back
		vec[i].push_back(1e9);
		for (int j = 0; j < vec[i].size()-1; j++) {
			// 자신의 등수보다 K 보다 큰 마지막 친구가 iter로 반환 됩니다.
			auto iter = upper_bound(vec[i].begin(), vec[i].end(), vec[i][j] + K); iter--;
			// 시작 주소를 빼줌으로써 index값을 얻을 수 있고
			auto index = iter - vec[i].begin();
			// 마지막 친구의 index - 현재 index를 해주면 친구의 명수가 나오게 됩니다.
			ans += index - j;
		}
	}
	cout << ans << "\n";
}