[SWEA 3947] 가장 짧은 길 전부 청소하기 - C++

SW Expert Academy

문제

석환이가 사는 동네에는 석환이 집을 포함해서 건물이 N개 있다.
건물은 1부터 N까지 번호가 붙어 있다. 석환이 집은 1번 건물이다.
또, 동네에는 M개의 길이 있는데, 각 길은 두 건물을 연결하고 양방향으로 통행이 가능하다.
석환이는 올림피아드 출제에 소홀했다는 이유로 벌청소를 하게 되었다.
석환이는 길들 중 일부를 청소해야 한다.
석환이는 자신이 이용하는 길만 청소하면 된다.
석환이가 길들을 이용하는 방법은 자기 집에서 어떤 한 건물로 가는 최단경로를 이용해서 그 건물에 갔다가 다시 집으로 돌아오는 것이다.
각 길을 청소하는 비용은 길의 길이와 같은 값이다.
석환이는 동네에 있는 모든 건물을 방문할 수 있으므로, 각 건물에 대해서 석환이 집에서 그 건물에 가는 최단경로에 존재하는 길들을 청소해야 한다.
한 건물에 대해서, 석환이 집에서 그 건물로 가는 최단 경로는 여러 개가 존재할 수 있는데, 그 중 하나의 최단 경로에 있는 길들만 청소하면 된다.
아래의 예를 보면 1번 건물(석환이 집)에서 4번 건물로 갈수 있는 최단 경로는 두가지가 존재하고 5번 건물로 갈수 있는 최단 경로도 두 가지가 존재한다.
4번과 5번으로 가는 최단 경로들을 모두 2번 건물을 거쳐 가도록 정하면 최소 청소 비용 7을 얻는다.
만약, 4번 건물로 가는 최단 경로는 2번 건물을 거치도록 정하고, 5번 건물로 가는 최단 경로는 3번 건물을 거쳐가도록 정하는 경우 청소 비용은 9가 되어 최소가 아니다.

석환이가 사는 동네의 길들이 건물들을 연결하는 정보와 각 길의 길이가 주어졌을 때, 석환이의 청소 비용의 최소값을 계산하는 프로그램을 작성하라.

입력

첫 줄에 테스트케이스의 개수 T가 주어진다. (T ≤ 20)
각 테스트케이스의 첫째 줄에 건물의 수를 나타내는 자연수 N과 길들의 수를 나타내는 자연수 M이 주어진다 (1 ≤ N ≤ 200,000, 0 ≤ M ≤ 500,000).
다음 이어지는 M개의 줄 각각에는 하나의 길을 나타내는 3개의 자연수가 주어진다.
그 중 첫 두 수는 서로 다른 건물의 번호이고 세 번째 수는 그 길의 길이 이다.
각 길의 길이는 1이상 10억 이하이다.
동일한 쌍의 건물을 잇는 길은 최대 하나이다.
모든 건물 사이에 길들을 이용해서 이동할 수 있는 경로가 존재한다.

풀이

#include<iostream>
#include<vector>
#include<queue>
#include<limits>
#include<algorithm>
#include<cstring>
 
using namespace std;
 
const long long INF = numeric_limits<long long>::max();
const int UNVISITED = -1;
 
struct Node {
    int ID;
    long long cost;
    bool checked;
    Node(int ID, long long cost) : ID(ID), cost(cost), checked(false) {}
    bool operator<(const Node& other) const {
        return cost > other.cost; // 최소 힙을 위해 비용이 작은 순서로 정렬
    }
};
 
int N, M;
long long answer;
vector<vector<Node>> adList;
vector<long long> dist;
vector<int> trace;
vector<bool> visited;
 
Node* findEdge(int A, int B) {
    for (auto& edge : adList[A]) {
        if (edge.ID == B) {
            return &edge;
        }
    }
    return nullptr;
}
 
void traceReturn(int start) {
    for (int v = 1; v <= N; v++) {
        if (v == start) continue;
 
        int prev = v;
        while (prev != start) {
            int A = prev;
            int B = trace[prev];
 
            Node* edge = findEdge(A, B);
            if (edge->checked) break;
 
            answer += edge->cost;
            edge->checked = true;
 
            Node* converseEdge = findEdge(B, A);
            if (converseEdge) converseEdge->checked = true;
 
            prev = trace[prev];
        }
    }
}
 
void dijkstra(int start) {
    priority_queue<Node> pq;
    dist[start] = 0;
    trace[start] = start;
    pq.push(Node(start, 0));
 
    while (!pq.empty()) {
        Node current = pq.top();
        pq.pop();
 
        if (visited[current.ID]) continue;
        visited[current.ID] = true;
 
        for (auto& next : adList[current.ID]) {
            if (dist[next.ID] >= dist[current.ID] + next.cost) {
                dist[next.ID] = dist[current.ID] + next.cost;
 
                if (trace[next.ID] == UNVISITED) {
                    trace[next.ID] = current.ID;
                }
                else {
                    Node* A = findEdge(current.ID, next.ID);
                    Node* B = findEdge(trace[next.ID], next.ID);
                    if (A && B && A->cost < B->cost) {
                        trace[next.ID] = current.ID;
                    }
                }
 
                pq.push(Node(next.ID, dist[next.ID]));
            }
        }
    }
    traceReturn(start);
}
 
int main(int argc, char** argv)
{
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    cout.tie(NULL);
    int test_case;
    int T;
    //freopen("input.txt", "r", stdin);
    cin >> T;
 
    for (test_case = 1; test_case <= T; ++test_case)
    {
        cin >> N >> M;
        adList.assign(N + 1, vector<Node>());
        dist.assign(N + 1, INF);
        trace.assign(N + 1, UNVISITED);
        visited.assign(N + 1, false);
 
        for (int i = 0; i < M; i++) {
            int u, v;
            long long cost;
            cin >> u >> v >> cost;
 
            adList[u].emplace_back(v, cost);
            adList[v].emplace_back(u, cost);
        }
        answer = 0;
        dijkstra(1);
        cout << "#" << test_case << " " << answer << "\n";
    }
    return 0;//정상종료시 반드시 0을 리턴해야합니다.
}

해결방법

문제의 예제로 예상되는 최소 신장 트리의 프림 알고리즘으로는 해결할 수 없는 문제입니다. 해당 문제를 해결하는 순서는 다음과 같습니다.

1. 그래프 정보를 입력받는다.
2. 다익스트라 알고리즘을 실행하여 최단 거리 찾기.
3. 특정 경로를 추적하여 결과 값을 계산하고 출력.

해당 문제는 시간이 굉장히 타이트한 문제여서 코드에서 확인할 수 있듯이 입출력 관련한 추가 설정을 통해 통과하였습니다. 위 코드의 시간 복잡도를 더 줄일 수 있는 아이디어는 다음과 같습니다.

1. findEdge() 호출 비용을 줄이기위해 unordered_map<int, Node>로 변경하기.
2. visited 배열로 중복 방문을 관리해서 checked 변수 제거하기.
3. trace 배열을 활용한 역추적 방법 개선하기.