Softeer-연습문제-지우는 소수를 좋아해(C++)

문제 링크 :

https://softeer.ai/practice/info.do?eventIdx=1&psProblemId=582

 

문제 개요 : 

마지막 체육관을 가기 위한 최소 소수 찾기

 

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

#define INF 1987654321

using namespace std;

int N,M;
vector<int> prime_num; // 소수 저장
vector<vector<pair<int,int>>> map; // 길 저장
int find_prime_num(int); // 소수 찾는 함수
int solution(); // 최소 레벨 찾는 함수

int main(int argc, char** argv)
{
	// 입력
	cin >> N >> M;
	map.resize(N);
	for(int i=0;i<M;i++){
		int A,B,C;
		cin >> A >> B >> C;
		map[--A].push_back(make_pair(--B,C));
		map[B].push_back(make_pair(A,C));
	}

	// 소수 계산
	prime_num.push_back(2);
	prime_num.push_back(3);
	prime_num.push_back(5);
	int tmp=prime_num.back()+2;
	while(tmp*tmp<=INF){
		bool is_prime = 1;
		for(int j=0;j<prime_num.size();j++){
			if(tmp % prime_num[j] != 0) continue;
			is_prime=0;
			break;
		}
		if(is_prime) prime_num.push_back(tmp); // 소수 찾음
		tmp+=2; // 홀수만 소수인지 확인
	}

	cout << solution();

	return 0;
}

int solution(){
	vector<int> Lv(N,INF);
	queue<pair<int,int>> q;
	for(int i=0;i<map[0].size();i++){ // 큐 초기화
		int des=map[0][i].first,lv=map[0][i].second;
		q.push(make_pair(des,lv));
		Lv[des]=lv;
	}
	while(!q.empty()){
		int idx = q.front().first; // 현재 체육관
		int tmp_lv = q.front().second; // 현재 레벨
		q.pop();
		for(int i=0;i<map[idx].size();i++){
			int des = map[idx][i].first; // 다음 체육관
			int lv = max(map[idx][i].second,tmp_lv); // 다음 레벨
			if(lv<Lv[des]){ // 더 작은 레벨을 Lv에 저장
				Lv[des]=lv;
				q.push(make_pair(des,lv));
			}
		}
	}
	return find_prime_num(Lv[N-1]+1); // Lv+1 해줘야 함...
}

int find_prime_num(int n){
	for(int i=0;i<prime_num.size();i++){
		if(prime_num[i]*prime_num[i]>n) break;
		if(n<prime_num[i]) return prime_num[i];
		if(n%prime_num[i]==0) return find_prime_num(n+1);
	}
	return n;
}

 

문제 해결 방법

1. 필자는 다음과 같은 형식으로 길을 저장했다. A체육관에서 B체육관까지 가는데 필요한 레벨 C를 <B,C>의 형식으로 map 벡터의 A행에 저장했다.

2. prime_num 벡터에 소수들을 작은 수부터 저장했다. 이때 소수는 문제를 풀기에 충분히 큰 수까지 찾아준다.

3. solution 함수에서는 각 체육관을 가기 위한 최소 레벨을 Lv 벡터에 저장한다.

3-1. 다음으로 방문할 체육관과 필요레벨을 큐에 저장하면 쉽게 최소레벨을 구할 수 있다.

3-2. 최소레벨을 구한 다음 +1을 해주어야 정답처리를 받을 수 있다. 이건 문제에서 제대로 알려주지 않았다. 필자는 이것 때문에 고생 좀 했다.

4. 앞에서 구한 최소레벨보다 크거나 같은 최소 소수를 찾는다. 이때 prime_num 에 저장된 소수를 이용한다.

5. 정답을 출력한다.