[백준] 1504번 특정한 최단 경로 : 다익스트라 - C++

---

1504번 특정한 최단 경로

 노드 n개가 주어질 때 시작점에서 v1, v2를 모두 경유하여 끝점으로 가는 최단경로를 구하는 문제입니다. 가능한 경우의 수는 두 가지뿐입니다. 첫 번째, 시작점 - v1 - v2 - 끝점. 두 번째, 시작점 - v2 - v1 - 끝점. 두 경로중 거리가 더 짧은 경로를 출력하면 됩니다. 시작점에서 v1, v2까지의 최단거리와 v1, v2사이의 최단거리, v1, v2에서 끝점사이의 최단거리를 구하면 됩니다.

풀이방법

 처음에는 학교에서 배운 다익스트라를 그대로 적용하여 distance배열을 업데이트 하는 방식을 사용했습니다. 하지만 distance배열에서 최솟값을 찾고 업데이트를 하는 과정이 오래 걸려 시간초과가 발생했습니다.

// 시간초과 코드

#include <iostream>
#include <vector>

using namespace std;

int main()
{
	// 입력
	int node_num, edge_num;
	cin >> node_num >> edge_num;

	vector<vector<int>> distance(node_num, vector<int>(node_num, 10000));
	for (int i = 0; i < edge_num; ++i)
	{
		int start, end, cost;
		cin >> start >> end >> cost;
		distance[start - 1][end - 1] = cost;
		distance[end - 1][start - 1] = cost;
	}

	int v1, v2;
	cin >> v1 >> v2;
	--v1;
	--v2;

	// 1
	{
		vector<bool> visited(node_num);
		visited[0] = true;
		while (!visited[v1] || !visited[v2])
		{
			int min_node;
			int min_dist = 10000;
			for (int i = 0; i < node_num; ++i)
				if (!visited[i] && distance[0][i] < min_dist)
				{
					min_node = i;
					min_dist = distance[0][i];
				}

			visited[min_node] = true;
			for (int i = 0; i < node_num; ++i)
				if (!visited[i] && min_dist + distance[min_node][i] < distance[0][i])
					distance[0][i] = min_dist + distance[min_node][i];
		}
	}

	// v1
	{
		vector<bool> visited(node_num);
		visited[v1] = true;
		while (!visited[node_num - 1] || !visited[v2])
		{
			int min_node;
			int min_dist = 10000;
			for (int i = 0; i < node_num; ++i)
				if (!visited[i] && distance[v1][i] < min_dist)
				{
					min_node = i;
					min_dist = distance[v1][i];
				}

			visited[min_node] = true;
			for (int i = 0; i < node_num; ++i)
				if (!visited[i] && min_dist + distance[min_node][i] < distance[v1][i])
					distance[v1][i] = min_dist + distance[min_node][i];
		}
	}

	// v2
	{
		vector<bool> visited(node_num);
		visited[v2] = true;
		while (!visited[node_num - 1])
		{
			int min_node;
			int min_dist = 10000;
			for (int i = 0; i < node_num; ++i)
				if (!visited[i] && distance[v2][i] < min_dist)
				{
					min_node = i;
					min_dist = distance[v2][i];
				}

			visited[min_node] = true;
			for (int i = 0; i < node_num; ++i)
				if (!visited[i] && min_dist + distance[min_node][i] < distance[v2][i])
					distance[v2][i] = min_dist + distance[min_node][i];
		}
	}

	cout << min(distance[0][v1] + distance[v1][v2] + distance[v2][node_num - 1],
		distance[0][v2] + distance[v1][v2] + distance[v1][node_num - 1]);

	return 0;
}

 

 그래서 우선순위큐를 사용해보았습니다. 우선순위 큐를 활용하면 최소 간선을 빠르게 선택할 수 있습니다.

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

using namespace std;

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

int main()
{
	int node_num, edge_num;
	cin >> node_num >> edge_num;

	vector<vector<int>> distance(node_num, vector<int>(node_num, -1));
	for (int i = 0; i < edge_num; ++i)
	{
		int start, end, cost;
		cin >> start >> end >> cost;
		distance[start - 1][end - 1] = cost;
		distance[end - 1][start - 1] = cost;
	}

	int v1, v2;
	cin >> v1 >> v2;
	--v1;
	--v2;

	// 1 to v1, v2
	int startToV1 = -1;
	int startToV2 = -1;
	{
		priority_queue<pair<int, int>, vector<pair<int, int>>, Comp> bfs;
		vector<bool> visited(node_num, false);
		bfs.push({ 0,0 });
		while (!bfs.empty() && (!visited[v1] || !visited[v2]))
		{
			int curr = bfs.top().first;
			int cost = bfs.top().second;
			bfs.pop();

			if (visited[curr])
				continue;
			visited[curr] = true;

			if (curr == v1)
				startToV1 = cost;
			if (curr == v2)
				startToV2 = cost;

			for (int i = 0; i < node_num; ++i)
				if (!visited[i] && distance[curr][i] != -1)
					bfs.push({ i,cost + distance[curr][i] });
		}
	}

	// v1 to v2, n-1
	int v1ToV2 = -1;
	int v1ToEnd = -1;
	{
		priority_queue<pair<int, int>, vector<pair<int, int>>, Comp> bfs;
		vector<bool> visited(node_num, false);
		bfs.push({ v1,0 });
		while (!bfs.empty() && (!visited[v2] || !visited[node_num - 1]))
		{
			int curr = bfs.top().first;
			int cost = bfs.top().second;
			bfs.pop();

			if (visited[curr])
				continue;
			visited[curr] = true;

			if (curr == v2)
				v1ToV2 = cost;
			if (curr == node_num - 1)
				v1ToEnd = cost;

			for (int i = 0; i < node_num; ++i)
				if (!visited[i] && distance[curr][i] != -1)
					bfs.push({ i,cost + distance[curr][i] });
		}
	}

	// v2 to n-1
	int v2ToEnd = -1;
	{
		priority_queue<pair<int, int>, vector<pair<int, int>>, Comp> bfs;
		vector<bool> visited(node_num, false);
		bfs.push({ v2,0 });
		while (!bfs.empty() && !visited[node_num - 1])
		{
			int curr = bfs.top().first;
			int cost = bfs.top().second;
			bfs.pop();

			if (visited[curr])
				continue;
			visited[curr] = true;

			if (curr == node_num - 1)
				v2ToEnd = cost;

			for (int i = 0; i < node_num; ++i)
				if (!visited[i] && distance[curr][i] != -1)
					bfs.push({ i,cost + distance[curr][i] });
		}
	}

	int path1, path2;
	if (startToV1 == -1 || v1ToV2 == -1 || v2ToEnd == -1)
		path1 = -1;
	else
		path1 = startToV1 + v1ToV2 + v2ToEnd;

	if (startToV2 == -1 || v1ToV2 == -1 || v1ToEnd == -1)
		path2 = -1;
	else
		path2 = startToV2 + v1ToV2 + v1ToEnd;

	if (path1 == -1 && path2 == -1)
		cout << -1;
	else if (path1 == -1)
		cout << path2;
	else if (path2 == -1)
		cout << path1;
	else
		cout << (path1 > path2 ? path2 : path1);

	return 0;
}

 

마치며

 처음 문제를 보고 쉽게 풀 수 있을 줄 알았는데 시간초과가 떠서 생각보다 오래 걸렸습니다. 아직 시간복잡도를 생각하는 것이 익숙지 않은 것 같습니다. 그래도 우선순위큐를 사용해야겠다고 생각했을 때 전에 비슷한 문제인 숨바꼭질 3 을 풀어봤어서 금방 코드를 짠 거 같습니다. 그래도 경험치가 조금씩 쌓이는 게 느껴지니 꽤 뿌듯합니다.