백준 5719번 거의 최단 경로 (C++)

문제링크

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

설명

정답률이 19%라 잔뜩 쫄았는데 시킨대로 구현만 잘하니 바로 AC가 나와 놀란 문제다.
전체적인 풀이 순서는 다음과 같다.

1. 다익스트라로 도착지까지의 최단 경로 값을 구해놓는다.
2. 최단 경로를 역추적하여 사용한 경로를 기록해둔다.
3. 사용한 경로를 사용하지 않고 다익스트라를 다시 실행한다.

테스트 케이스가 여러 개이기도 하고, 다익스트라를 두 번 실행하기 때문에 이전에 사용한 값들을 잘 초기화시켜주어야 한다.

코드는 아래와 같다.

코드

#include <iostream>
#include <vector>
#include <queue>
#include <cstring>
#define N 505
#define M 10010
#define p pair<int, int>

using namespace std;

struct node {
	int next, dis, idx;
};

int n, m, s, d, dis[N], val[N][N];	// N의 최댓값이 500이므로, cost를 기록해 줄 수 있다.
vector<node> con[N];
vector<p> pre[N];	// a -> b이고 c번째 경로일 때, pre[b]에 a와 c를 기록
bool check[M];	// i번째 경로 사용 기록

void reset_dis()	// 각각의 다익스트라 때 초기화해준다.
{
	for (int i = 0; i < n; i++)
		dis[i] = 1e9;
}

void reset_rest()	// 다음 테케를 위해 사용한 것들을 초기화해준다.
{
	for (int i = 0; i < n; i++)
	{
		con[i].clear();
		pre[i].clear();
	}

	memset(check, 0, sizeof(check));
}

void input()
{
	cin >> s >> d;

	for (int i = 0; i < m; i++)
	{
		int a, b, c;
		cin >> a >> b >> c;
		con[a].push_back({ b, c, i });
		pre[b].push_back({ a, i });
		val[a][b] = c;
	}
}

void init()	// 첫 다익스트라 전처리
{
	priority_queue<p> pq;
	pq.push({ 0, s });
	dis[s] = 0;

	while (!pq.empty())
	{
		int cost = pq.top().first * -1;
		int cur = pq.top().second;
		pq.pop();

		if (dis[cur] < cost || cur == d)
			continue;

		for (auto i : con[cur])
		{
			int next = i.next;
			int ncost = cost + i.dis;

			if (dis[next] <= ncost)
				continue;

			pq.push({ ncost * -1, next });
			dis[next] = ncost;
		}
	}
}

// dfs를 이용한 경로 역추적
// 이 때 모든 경로를 지워주어야하기 때문에 조건을 만족한 모든 경로를 지워주는 것을 잊지말자
void trace(int cur)	
{
	if (cur == s)
		return;

	for (int i = 0; i < pre[cur].size(); i++)
	{
		p temp = pre[cur][i];

		if (dis[cur] != dis[temp.first] + val[temp.first][cur] || check[temp.second])
			continue;

		check[temp.second] = 1;
		trace(temp.first);
	}
}

void solution()
{
	init();
	trace(d);
	reset_dis();

	// 두 번째 다익스트라
	int ans = -1;
	priority_queue<p> pq;
	pq.push({ 0, s });
	dis[s] = 0;

	while (!pq.empty())
	{
		int cost = pq.top().first * -1;
		int cur = pq.top().second;
		pq.pop();

		if (dis[cur] < cost)
			continue;

		if (cur == d)
		{
			ans = cost;
			break;
		}

		for (auto i : con[cur])
		{
			int next = i.next;
			int ncost = cost + i.dis;
			int idx = i.idx;

			if (dis[next] <= ncost || check[idx])
				continue;

			pq.push({ ncost * -1, next });
			dis[next] = ncost;
		}
	}

	cout << ans << '\n';
}

void solve()
{
	while (1)
	{
		cin >> n >> m;

		if (!n && !m)
			break;

		reset_dis();
		input();
		solution();
		reset_rest();
	}
}

int main(void)
{
	ios_base::sync_with_stdio(false);
	cin.tie(NULL);

	solve();
	return 0;
}