백준 23793번 두 단계 최단 경로 1 C++ 다익스트라

리뷰

특정 노드를 무조건 방문하고, 방문하지 않는 조건을 가진 최단 경로 문제

 

문제 풀이

1. n, m을 입력 받고 인접 리스트를 n + 1 크기로 초기화 해준 뒤 단방향으로 노드와 가중치를 벡터에 넣어준다.
2. x, y, z를 입력 받고 first에 x 에서 y까지의 최단 경로를 구하고, y부터 z까지의 최단 경로를 구한 뒤 더해준다.
3. second에 x부터 z까지 최단 경로를 구해준다. 이때 bfs에 flag를 전달 해 y를 들릴지 말지를 정해준다.
4. bfs가 실행되면 n + 1크기의 거리 배열을 매우 큰 값으로 초기화 해준 뒤 최소힙을 만들어 준다.
5. 시작 위치를 힙에 넣고, 거리를 0으로 초기화 해준 뒤 while 루프를 돌려준다.
6. 만약 bfs에 전달된 flag가 0이라면 기존 bfs를 진행, flag가 1이라면 다음 노드가 y일때 continue 처리를 해준다.
7. while루프를 끝내고 목적지 까지의 거리를 리턴해준다.
8. first와 second 모두 만약 출력된 값이 INF 보다 크다면 목적지에 도달하지 못한 것이므로 -1을 출력
9. 아니라면 first와 second를 각각 출력해 주면 된다.

 

참고 사항

가중치는 최대 10000이고, m값이 최대 200000이므로 최종 거리가 int를 넘을 수도 있을 것 같아 long long 타입으로 세팅해 주었다.

 

정답 코드

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

using namespace std;

const long long INF = 9999999999;
int n, m, x, y, z;
vector<vector<pair<long long, int>>> lst;

long long bfs(int start, int end, int flag) {
    vector<long long> dist(n + 1, INF);
    priority_queue<pair<long long, int>, vector<pair<long long, int>>, greater<>> pq;

    pq.push({ 0, start });
    dist[start] = 0;

    while (!pq.empty()) {
        pair<long long, int> cur = pq.top(); pq.pop();
        long long cd = cur.first;
        int cn = cur.second;
        if (dist[cn] != cd) continue;

        for (pair<long long, int> next : lst[cn]) {
            long long nd = next.first;
            int nn = next.second;
            if (flag && nn == y) continue;
            if (dist[nn] > cd + nd) {
                dist[nn] = cd + nd;
                pq.push({ dist[nn], nn });
            }
        }
    }
    return dist[end];
}

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

    cin >> n >> m;

    lst.resize(n + 1);
    for (int i = 0; i < m; i++) {
        int a, b, c;
        cin >> a >> b >> c;
        lst[a].push_back({ c, b });
    }
    
    cin >> x >> y >> z;
    long long first = bfs(x, y, 0) + bfs(y, z, 0);
    long long second = bfs(x, z, 1);
    
    if (first >= INF) cout << -1 << " ";
    else cout << first << " ";
    if (second >= INF) cout << -1;
    else cout << second;
}