백준 1854 K번째 최단경로 찾기

백준 1854←클릭

변수 설정

town_dist[k]: k번째 마을까지 가는 거리를 저장한 우선순위 큐. size는 k로 제한하고, 내림차순이기 때문에 top()은 k번째로 가까운 거리가 유지된다.
dist[from][to]:from에서 to까지 간선의 길이 저장
pq: 시작 노드에서 특정 노드까지의 거리를 pair형태로 저장하는 우선순위 큐. second값을 기준으로 오름차순이다.

아이디어

• 일반적으로 다익스트라는 시작점 기준 가장 가까운 노드를 거쳐 다른 노드를 갈 때 최소 거리를 갱신하면서 간다.
  
  
• 다만 해당 문제는 최소 거리를 갱신하는 것이 아닌 k번째 가까운 거리를 구해야 한다.
  
  
• 노드마다 우선순위 큐를 내림차순으로 만들고 시작 노드 → 중간 노드 + 중간 노드 →목적지 노드를 계산한 뒤 특정 조건을 만족하면 우선순위 큐(town_dist)에 넣는다.
  
  
• 특정 조건을 만족하는 경우 pq에도 새로 계산된 거리를 삽입하여 나중에 다시 방문한다.
  
  
  우선 순위 큐에 넣는 특정 조건은 다음과 같다.

우선 순위 큐 삽입 조건

• 아직 town_dist가 k만큼 안찬 경우
  
  다음과 같이 아직 k만큼 town_dist가 안차면 아묻따 넣는다. 우선 순위 큐 이기에 알아서 정렬이 된다.
  
  
• town_dist가 k만큼 찼지만, top()보다 작은 거리를 찾는 경우
  
  다음과 같이 top()보다 작은 거리를 찾는 경우에 우선 순위 큐를 최신화한다.

이 두 경우에 town_dist에 새로운 거리를 추가하는 것 이외에도 pq에도 추가하여 나중에 해당 거리로 같은 노드를 한번 더 방문한다. 이 문제는 노드를 한번만 방문하라는 조건이 없기 때문이다.

코드

github

#define _CRT_SECURE_NO_WARNINGS
#define INF 987654321
#include<iostream>
#include<algorithm>
#include<vector>
#include<queue>
using namespace std;

void setting();
void dijkstra();
void solution();
void answer();

bool print = false;
int n, m, k;
vector<vector<int>> dist(1001,vector<int>(1001,INF));
vector<priority_queue<int, vector<int> , less<int>>> town_dist(1001); //내림차순

struct cmp {
	bool operator() (pair<int, int>& a, pair<int, int>& b) {
		return a.second > b.second;
	}
};

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

	solution();

	return 0;
}

void setting() {
	if (print) cout << "setting" << endl;
	//freopen("input/1854_input.txt", "r", stdin);
	cin >> n >> m >> k;
	for (int i = 1; i <= m; i++) {
		int a, b, c;
		cin >> a >> b >> c;
		dist[a][b] = c;
	} 

	town_dist[1].push(0); //시작점->시작점의 최소값은 0
}

void solution() {
	setting();
	dijkstra();
	answer();
}
	
void dijkstra() {
	priority_queue<pair<int, int>, vector<pair<int, int>>, cmp> pq; //second기준 오름차순
	pq.push(make_pair(1, 0));

	while (!pq.empty()) {
		int current = pq.top().first; //현재 노드 번호
		int distance = pq.top().second; //거리
		pq.pop();
		//if (min_dist[current] < distance) continue; //최소가 아닌경우 continue
		if (print)printf("현재 노드: (%d, %d)\n", current, distance);

		for (int i = 1; i <= n; i++) {
			int next, next_dist;
			//인접 노드
			if (dist[current][i] == INF) continue; //간선이 없는 경우
			else {
				next = i; //인접 노드 번호
				next_dist = distance + dist[current][next];
			}
			if (print)printf("next: %d, next_dist:%d, size: %d\n", next, next_dist, town_dist.at(next).size());
			if (town_dist.at(next).size() < k) { //k만큼 아직 못채운 경우
				town_dist.at(next).push(next_dist);
				pq.push(make_pair(next, next_dist));//pq에 삽입
			}
			else if (town_dist.at(next).top() > next_dist) { //k만큼 채웠는데 갱신
				if (print) cout << " top 갱신 " << endl;
				town_dist.at(next).pop();
				town_dist.at(next).push(next_dist);
				pq.push(make_pair(next, next_dist));//pq에 삽입
			}
			// k만큼 채웠지만 갱신을 못하는 경우에는 pq에 다시 넣지 않는다. 
		}
	}
}

void answer() {
	for (int i = 1; i <= n; i++) {
		if(print)cout << "i: " << i << endl;
		if (k > town_dist[i].size()) cout << "-1" << endl;
		else cout << town_dist[i].top() << endl;
	}
}