문제 링크: https://algospot.com/judge/problem/read/FIRETRUCKS

 

다익스트라 알고리즘을 사용하면 되는데요, 머리속에 떠오르는 그대로 풀면 틀릴 가능성이 높습니다.

저는 딱 보자마자 '각 소방서마다 다익스트라 돌려서 경로중에 가장 짧은거 골라서 더하면 되겠네' 라고 생각했지만 그렇게 하면 O(mElgV)의 시간복잡도가 나오고 간선의 수의 최대치가 50만에 가깝기 때문에 불가능합니다.

 

책을 보니 무릎을 탁 치는 해법이 있는데요, 그래프에 가상의 시작점을 하나 추가한 뒤, 이 시작점에서 다른 모든 소방서로 가중치 0인 간선을 연결하는 것입니다.

이렇게 하면 다익스트라 한 번의 수행으로 모든 위치에 대해 임의의 소방서로부터의 최단 거리를 얻을 수 있습니다.

#pragma GCC optimize ("Ofast")

#define _CRT_SECURE_NO_WARNINGS
#include <bits/stdc++.h>
using namespace std;


int v, e, n, m;
vector<vector<pair<int, int>>> adj;	// idx, cost

vector<int> dijkstra(int src) {
	vector<int> dist(v, numeric_limits<int>::max());
	dist[src] = 0;

	priority_queue<pair<int, int>> pq;	// cost(-), idx
	pq.emplace(0, src);

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

		if (dist[here] < cost) continue;

		for (const auto& edge : adj[here]) {
			int there = edge.first;
			int nextDist = cost + edge.second;

			if (dist[there] > nextDist) {
				dist[there] = nextDist;
				pq.emplace(-nextDist, there);
			}
		}
	}

	return dist;
}

int main()
{
#ifndef ONLINE_JUDGE
	freopen("input.txt", "r", stdin);
#endif

	ios_base::sync_with_stdio(false);
	cin.tie(nullptr);
	cout.tie(nullptr);

	int t;
	cin >> t;
	while (t--) {
		adj.clear();
		cin >> v >> e >> n >> m;
		++v;

		adj.resize(v);
		while (e--) {
			int a, b, t;
			cin >> a >> b >> t;
			adj[a].emplace_back(b, t);
			adj[b].emplace_back(a, t);
		}

		vector<int> fire(n);
		for (auto& f : fire)
			cin >> f;

		while (m--) {
			int i;
			cin >> i;
			adj[0].emplace_back(i, 0);
		}

		auto dist = dijkstra(0);
		long long sum = 0;
		for (auto& f : fire)
			sum += dist[f];
		cout << sum << '\n';
	}

	return 0;
}
반응형