[프로그래머스/C++] 부대복귀

미숙한 블로그 주인이 코딩테스트 문제를 풀어가는 과정을 담은 글입니다. 이 풀이가 효율적인 풀이가 아닐 수 있으며, 부정확한 정보가 많이 있을 수 있습니다. 보완해야할 점이 있다면 댓글로 남겨주세요!

 

https://school.programmers.co.kr/learn/courses/30/lessons/132266

프로그래머스

부대복귀

문제

강철부대의 각 부대원이 여러 지역에 뿔뿔이 흩어져 특수 임무를 수행 중입니다. 지도에서 강철부대가 위치한 지역을 포함한 각 지역은 유일한 번호로 구분되며, 두 지역 간의 길을 통과하는 데 걸리는 시간은 모두 1로 동일합니다. 임무를 수행한 각 부대원은 지도 정보를 이용하여 최단시간에 부대로 복귀하고자 합니다. 다만 적군의 방해로 인해, 임무의 시작 때와 다르게 되돌아오는 경로가 없어져 복귀가 불가능한 부대원도 있을 수 있습니다.

강철부대가 위치한 지역을 포함한 총지역의 수 n, 두 지역을 왕복할 수 있는 길 정보를 담은 2차원 정수 배열 roads, 각 부대원이 위치한 서로 다른 지역들을 나타내는 정수 배열 sources, 강철부대의 지역 destination이 주어졌을 때, 주어진 sources의 원소 순서대로 강철부대로 복귀할 수 있는 최단시간을 담은 배열을 return하는 solution 함수를 완성해주세요. 복귀가 불가능한 경우 해당 부대원의 최단시간은 -1입니다.

 

제한사항

• 3 ≤ n ≤ 100,000
  • 각 지역은 정수 1부터 n까지의 번호로 구분됩니다.
• 2 ≤ roads의 길이 ≤ 500,000
  • roads의 원소의 길이 = 2
  • roads의 원소는 [a, b] 형태로 두 지역 a, b가 서로 왕복할 수 있음을 의미합니다.(1 ≤ a, b ≤ n, a ≠ b)
  • 동일한 정보가 중복해서 주어지지 않습니다.
    • 동일한 [a, b]가 중복해서 주어지지 않습니다.
    • [a, b]가 있다면 [b, a]는 주어지지 않습니다.
• 1 ≤ sources의 길이 ≤ 500
  • 1 ≤ sources[i] ≤ n
• 1 ≤ destination ≤ n

 

풀이

강철부대가 위치한 지역(sources)에서 목적지(destination)까지의 최단 거리를 구하는 문제라고 해석했다.

 

1차 풀이 (다익스트라 알고리즘)

#include <string>
#include <vector>
#include <queue>

#define INF 1999999999

using namespace std;

vector<int> solution(int n, vector<vector<int>> roads, vector<int> sources, int destination) {
    vector<int> answer;
    vector<vector<int>> maps(n + 1);

    for (vector<int> road : roads) {
        maps[road[0]].push_back(road[1]);
        maps[road[1]].push_back(road[0]);
    }

    for (int source : sources) {
        priority_queue<pair<int, int>> q;
        q.push({ 0, source });
        vector<bool> visited(n + 1, false);
        int result = INF;

        while (!q.empty()) {
            pair<int, int> p = q.top();
            int pos = p.second;
            int cost = -p.first;
            q.pop();

            if (cost >= result) {
                continue;
            }

            if (pos == destination) {
                result = min(result, cost);
                continue;
            }

            for (int i = 0; i < maps[pos].size(); i++) {
                if (!visited[maps[pos][i]]) {
                    q.push({ -(cost + 1), maps[pos][i] });
                    visited[maps[pos][i]] = true;
                }
            }
        }

        if (result == INF)
            answer.push_back(-1);
        else
            answer.push_back(result);
    }

    return answer;
}

그런데,, 이걸 통과한 것이라고 하기에는 좀 그렇다. 시간 제한이 조금만 더 엄격해지면 바로 시간 초과가 날 것이다..

그래서, 모든 지역들 사이의 거리가 1이라는 점을 이용해서, BFS로 풀이를 수정했다.

 

2차 풀이 (BFS)

#include <string>
#include <vector>
#include <queue>

#define INF 1999999999

using namespace std;

vector<int> solution(int n, vector<vector<int>> roads, vector<int> sources, int destination) {
    vector<int> answer;
    vector<vector<int>> maps(n + 1);

    for (vector<int> road : roads) {
        maps[road[0]].push_back(road[1]);
        maps[road[1]].push_back(road[0]);
    }

    for (int source : sources) {
        queue<pair<int, int>> q;
        q.push({ source, 0 });
        vector<bool> visited(n + 1, false);
        int result = INF;

        while (!q.empty()) {
            pair<int, int> p = q.front();
            int pos = p.first;
            int cost = p.second;
            q.pop();

            if (pos == destination) {
                result = cost;
                break;
            }

            for (int i = 0; i < maps[pos].size(); i++) {
                if (!visited[maps[pos][i]]) {
                    q.push({ maps[pos][i], cost + 1 });
                    visited[maps[pos][i]] = true;
                }
            }
        }

        if (result == INF)
            answer.push_back(-1);
        else
            answer.push_back(result);
    }

    return answer;
}

이것도 마음에 안 든다. 사실 두 풀이는 크게 차이가 없기도 하고,,

 

더 좋은 방법은 없을까 싶어 다른 분들의 풀이를 보던 중 한 분의 일침으로 엄청난 사실을 깨달았다..

목적지는 하나로 고정인데 sources에서 목적지까지의 거리를 구하고 있냐고.. 목적지에서 각 sources까지의 거리를 구하면 되는것 아니냐...는 것이었다.

 

나는 연산 한 번으로 끝나는 것을 여러번 돌렸던 것이었다.

 

최종 풀이 (destination에서 sources까지의 최단거리 구하기)

#include <string>
#include <vector>
#include <queue>
#include <map>

#define INF 1999999999

using namespace std;

vector<int> solution(int n, vector<vector<int>> roads, vector<int> sources, int destination) {
    vector<int> answer;
    vector<vector<int>> maps(n + 1);

    for (vector<int> road : roads) {
        maps[road[0]].push_back(road[1]);
        maps[road[1]].push_back(road[0]);
    }
    
    // 목적지에서 강철 부대의 각 지점들까지의 거리 저장용
    map<int, int> s;
    for (int source : sources) {
        s[source] = INF;
    }

    queue<pair<int, int>> q;
    vector<bool> visited(n + 1, false);

    q.push({ destination, 0 });
    visited[destination] = true;

    while (!q.empty()) {
        pair<int, int> p = q.front();
        int pos = p.first;
        int cost = p.second;
        q.pop();

        // 목적지에서 강철 부대의 위치로 도달했으면 거리 저장
        if (s[pos]) {
            s[pos] = cost;
        }
        
        for (int i = 0; i < maps[pos].size(); i++) {
            if (!visited[maps[pos][i]]) {
                q.push({ maps[pos][i], cost + 1 });
                visited[maps[pos][i]] = true;
            }
        }
    }

    for (int source : sources) {
        if (s[source] == INF)
            answer.push_back(-1);
        else
            answer.push_back(s[source]);
    }

    return answer;
}

속도가 훨 빨라졌다!