[백준 2606] 바이러스 with Python

문제 링크

www.acmicpc.net/problem/2606

2606번: 바이러스

풀이

import sys

class MyGraph():
    def __init__(self, n):
        self.graph = {}
        for i in range(1, n+1):
            self.graph[i] = []
    
    def add(self, start, end):
        self.graph[start].append(end)
        self.graph[end].append(start)

    def dfs(self, start):
        visited = []
        stack = [start]
        while stack:
            cur_node = stack.pop()
            if cur_node not in visited:
                visited.append(cur_node)
                stack.extend(self.graph[cur_node][::-1])
        return visited

n = int(sys.stdin.readline())
my_graph = MyGraph(n)

for _ in range(int(sys.stdin.readline())):
    start, end = map(int, sys.stdin.readline().split())
    my_graph.add(start, end)

print(len(my_graph.dfs(1))-1)

DFS로 해결하였다. 이 문제는 단방향 그래프가 아닌 양방향 그래프(무방향 그래프)라는 것에 유의해야 한다.

 

BFS로도 다시 풀어보았다.

import sys

class MyGraph():
    def __init__(self, n):
        self.graph = {}
        for i in range(1, n+1):
            self.graph[i] = []
    
    def add(self, start, end):
        self.graph[start].append(end)
        self.graph[end].append(start)

    def bfs(self, start):
        visited = [start]
        queue = [start]
        while queue:
            cur_node = queue.pop(0)
            for i in self.graph[cur_node]:
                if i not in visited:
                    queue.append(i)
                    visited.append(i)
        return visited

n = int(sys.stdin.readline())
my_graph = MyGraph(n)

for _ in range(int(sys.stdin.readline())):
    start, end = map(int, sys.stdin.readline().split())
    my_graph.add(start, end)

print(len(my_graph.bfs(1))-1)