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

프로그래머스

 

난이도 : Level 3 
알고리즘 유형 : BFS 
풀이 시간 : 17분 

문제 풀이 

각 부대원들이 여러 지역에서 임무를 수행중이다. 임무가 끝나고 부대로 복귀하는데 걸리는 최단 시간을 부대원 별로 구해서 출력하자. 

 

문제의 제한사항을 확인해보면 꽤나 빠듯하다. 먼저, [전체 정점 ≤ $100000$], [간선(roads)의 개수 ≤ $500000$], [병사의 수 ≤ $500$] 으로 주어졌기 때문에, 각 부대원 별로 한 명 씩 최단거리를 찾아주게 된다면 $O(500 * (100000 + 500000))$으로 시간이 오래 걸리게 된다.

 

따라서, 부대원 별로 거리를 탐색하는 것이 아닌, 반대로 생각해서, 목적지부터 각 정점까지의 최단거리를 탐색할 것이다. 

결국, 부대원들은 부대로 복귀를 해야하기 때문에, 목적지에서 각 정점으로 탐색을 시작해도 무방하다고 생각했다. 이렇게 탐색할 경우, 시간복잡도를 $O(100000 + 500000 + 500)$로 줄일 수 있기 때문에 제한시간 내에 안전하게 탐색이 가능하다.

 

따라서, 다음과 같은 순서로 풀이를 진행했다.

1. 양방향 그래프 초기화

2. 부대(`destination`)를 시작 정점으로 하여 모든 정점까지 BFS 진행

    2-1. visited 배열을 -1로 초기화 후, 탐색해나가며 부대 - 정점 간 최단거리를 저장

3. 각 부대원들의 위치(`source`)를 순회하며 `answer`에 visited[soldier] 값 push

 

코드

/*
각 지역은 유일한 번호로 구분
두 지역 간 길 통과 시간은 1로 동일
최단시간에 부대로 복귀하고자 한다.
적의 방해로 경로가 없어져 복귀가 불가능한 부대원이 있을 수도 있음

풀이 : visited를 -1로 설정한 뒤, destination을 시작으로 해서 각 노드까지 거리를 돌기
*/

#include <string>
#include <vector>
#include <iostream>
#include <cstring>
#include <queue>

using namespace std;

vector<vector<int>> graph(100001);
int visited[100001];

void bfs(int su){
    queue<int> q;
    visited[su] = 0;
    q.push(su);
    
    while (!q.empty()){
        int u = q.front();
        q.pop();
        
        for (auto v: graph[u]){
            if (visited[v] == -1){
                q.push(v);
                visited[v] = visited[u] + 1;
            }
        }
    }
}


vector<int> solution(int n, vector<vector<int>> roads, vector<int> sources, int destination) {
    vector<int> answer;
    memset(visited, -1, sizeof(visited));
    for (int i=0;i<roads.size();i++){
        graph[roads[i][0]].push_back(roads[i][1]);
        graph[roads[i][1]].push_back(roads[i][0]);
    }
    
    bfs(destination);
    
    for (auto loc : sources){
        answer.push_back(visited[loc]);
    }
    
    return answer;
}

배운 점 

사실, 부대원 각각에 대해 bfs를 수행해도 최대 수행시간이 1억이 넘지 않기 때문에(약 6천만) 실행은 가능할 것으로 생각했다. 하지만 시간을 더 최적화 할 수 있을 것이라고 생각하였고, 빠른 시간 내에 아이디어를 찾아서 진행하였다.

항상 최적화하는 방법을 고민하고 생각해나가는 습관을 가져야겠다.