백준 1504 특정한 최단 경로 (C++)

1504번: 특정한 최단 경로

다익스트라를 이용한 문제이다. 기존 다익스트라와 다른 점은 주어진 두 정점 v1, v2를 반드시 지나는 1부터 N까지의 최단 경로를 구하는 것이다. 이것은 두가지의 경우로 볼 수 있다.

1 -> v1 -> v2 -> N
1 -> v2 -> v1 -> N

아래 코드의 a, b, c가 각각 1, v1, N에서 각 정점까지의 최단 거리이다. 각 최단 거리를 더해 최소값을 구해서 출력해주었다.

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

#define INF 987654321

using namespace std;

int N, E, v1, v2, res = 0;
vector<pair<int, int>> v[801];


vector<int> dijkstra(int n) {
    priority_queue<pair<int, int>> pq;
    vector<int> d;

    d.push_back(0);
    for (int i = 1; i <= N; i++) {
        d.push_back(INF);
    }

    pq.push({ n,0 });
    d[n] = 0;

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

        pq.pop();

        for (int i = 0; i < v[cur].size(); i++) {
            int next = v[cur][i].first;
            int nc = cost + v[cur][i].second;

            if (nc < d[next]) {
                d[next] = nc;
                pq.push({ next,-nc });
            }
        }
    }

    return d;
}

void solution() {
    vector<int> a = dijkstra(1);
    vector<int> b = dijkstra(v1);
    vector<int> c = dijkstra(N);

    if (a[v1] == INF || b[v2] == INF || c[v2] == INF || a[v2] == INF || c[v1] == INF) res = -1;
    else res = min(a[v1] + b[v2] + c[v2], a[v2] + b[v2] + c[v1]);

    cout << res;
}

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

    cin >> N >> E;

    for (int i = 0; i < E; i++) {
        int a, b, c;
        cin >> a >> b >> c;
        v[a].push_back({ b,c });
        v[b].push_back({ a,c });
    }

    cin >> v1 >> v2;

    solution();

    return 0;
}