바이러스

2606번: 바이러스
https://www.acmicpc.net/problem/2606

문제

신종 바이러스인 웜 바이러스는 네트워크를 통해 전파된다. 한 컴퓨터가 웜 바이러스에 걸리면 그 컴퓨터와 네트워크 상에서 연결되어 있는 모든 컴퓨터는 웜 바이러스에 걸리게 된다.

예를 들어 7대의 컴퓨터가 <그림 1>과 같이 네트워크 상에서 연결되어 있다고 하자. 1번 컴퓨터가 웜 바이러스에 걸리면 웜 바이러스는 2번과 5번 컴퓨터를 거쳐 3번과 6번 컴퓨터까지 전파되어 2, 3, 5, 6 네 대의 컴퓨터는 웜 바이러스에 걸리게 된다. 하지만 4번과 7번 컴퓨터는 1번 컴퓨터와 네트워크상에서 연결되어 있지 않기 때문에 영향을 받지 않는다.

어느 날 1번 컴퓨터가 웜 바이러스에 걸렸다. 컴퓨터의 수와 네트워크 상에서 서로 연결되어 있는 정보가 주어질 때, 1번 컴퓨터를 통해 웜 바이러스에 걸리게 되는 컴퓨터의 수를 출력하는 프로그램을 작성하시오.

예제 입출력

// 입력				//출력
7    				4
6
1 2
2 3
1 5
5 2
5 6
4 7

해결 방향

1번 컴퓨터와 연결되어 있는 모든 컴퓨터의 수를 알아내야 한다.

깊이 우선 탐색(DFS)을 Stack으로 구현하여 1번 컴퓨터와 연결되어 있는 모든 컴퓨터를 방문하고, Set에 삽입한다.

코드

#include <iostream>
#include <vector>
#include <stack>
#include <set>
using namespace std;

int main() {
    int N, M, V=1, a, b, tmp=0;
    cin >> N >> M;
    vector<int> g[N+1];
    stack<int> st;
    set<int> s;
	
    // 각각의 컴퓨터에 연결되어 있는 다른 컴퓨터들의 번호를 저장한다. 
    for(int i=0; i<M; i++) {
        cin >> a >> b;
        g[a].push_back(b);
        g[b].push_back(a);
    }

	// 1번 컴퓨터를 stack과 set에 넣는다.
    st.push(1);
    s.insert(1);

    while(1) {
    	
        /* 어떤 컴퓨터가 set에 존재하지 않으면(방문한 적이 없으면),
           stack과 set에 넣는다. */
        if(s.find(tmp)==s.end() && tmp!=0) {
            V = tmp;
            st.push(V);
            s.insert(st.top());    
        }
        
        /* 현재 방문 중인 컴퓨터에 연결된 다른 컴퓨터를 이미 모두 방문했다면, 
		   현재 방문 중인 컴퓨터를 stack에서 pop하고, 이전에 방문했던 컴퓨터를 가져온다. */   
        if(g[V].size()==0) {
            st.pop();
            // stack이 비었다면, while문을 빠져나간다.
            if(st.empty()) break;
            V=st.top();
            continue;   
        }
        
        // 어떤 컴퓨터에 연결된 다른 컴퓨터를 가져온다.
        tmp = g[V].back();
        
        // 방문한 컴퓨터는 삭제한다.
        g[V].pop_back();
    }
	
    // 1번 컴퓨터를 제외한 바이러스에 전염된 컴퓨터의 개수를 출력한다.
    cout << s.size()-1;
}

회고

처음으로 깊이 우선 탐색 알고리즘을 사용하여 문제를 풀어보았다.
내가 구현한 코드는 가독성이 떨어지고, 메모리를 많이 사용하는 것 같다.
그래프 탐색 알고리즘에 대하여 더 공부해 보고, 코드를 개선해야겠다.