백준 2606 파이썬 | 바이러스

코드

from sys import stdin
input = stdin.readline

c = int(input())
l = int(input())
lst =[]
for i in range(l):
    a, b = map(int, input().split())
    lst.append([a, b])

nlst = set({1})
mlst = set()
olst = set()
while len(nlst) != len(olst) :
    for k in nlst:
        olst.add(k)
    for j in nlst:
        for i in range(l):
            if lst[i][1] == j or lst[i][0] == j:
                if lst[i][0] == j:
                    mlst.add(lst[i][1])
                else:
                    mlst.add(lst[i][0])
    nlst.update(mlst)
print(len(nlst)-1)

설명

1번 컴퓨터와 연결된 모든 다른 컴퓨터의 개수를 구하면 되는 프로그램이다. 입력으로 연결된 컴퓨터 번호 쌍이 주어진다. 탐색 알고리즘을 사용하는 것이 맞겠지만 나는 탐색 알고리즘을 몰라서 그냥 맘대로 반복문과 set 자료형을 이용하여 풀어 보았다. 

 

nlst 에는 1번 컴퓨터와 연결된 컴퓨터가 저장된다. 코드의 흐름은 while 문이 한 번 돌 때마다 nlst에 있는 컴퓨터와 연결된 컴퓨터가 있는지 입력 받은 컴퓨터 번호쌍을 이용해 하나하나 찾아준다. set 자료형을 쓰고 있어서 중복은 알아서 제거해 주기 때문에 막 추가해도 중복 문제는 걱정할 필요 없다. nlst에 있는 컴퓨터 번호가 입력 받은 컴퓨터 숫자쌍의 첫번째 자리에 있을지 두번째 자리에 있을지 모르기 때문에 그 경우를 대비해서 for i in range(l): 구문 안의 조건문을 짜주었다. 새롭게 추가되는 바이러스에 걸린 컴퓨터는 mlst에 추가한다. while 반복문의 마지막에 nlst와 mlst를 합쳐 준다. olst 는 while 반복문이 시작할 때 모든 nlst 내의 아이템을 복사해 온다. 만약 while 반목문을 한 번 돌고 다음 while 반복문이 시작할 때 nlst의 길이와 olst 의 길이가 같다면, 즉 nlst의 내용물이 변하지 않았다면 while 반복문을 탈출 할 수 있게 해 놓았다.

 

예제 입력을 통해 설명을 해 보겠다. 1번 컴퓨터가 바이러스의 시작이니까 nlst에는 이미 1이 들어있다. while 반복문을 첫번째로 돌면서 1번과 연결된 컴퓨터를 찾는다. (1, 2), (1, 5)가 찾아지고 nlst에 2, 5가 추가된다. while 반복문을 두번째로 돌 때, nlst 안에는 1, 2, 5가 있다. 각각의 번호에 대해 다시 반복문을 돌며 연결된 컴퓨터가 있는지 찾는다. 이 경우에는 (1, 2), (2, 3), (1, 5), (5, 2), (5, 6) 이 찾아진다. 중복은 set 자료형의 특성상 알아서 사라지고 다시 nlst에는 1, 2, 3, ,5, 6이 담긴다. while 반복문을 세번째로 돌 때, 더 이상 새롭게 연결된 컴퓨터 번호를 찾을 수가 없다. 그러면 while 조건에 의해서 nlst와 olst 의 길이가 동일하기 때문에 while 반복문을 탈출한다.

 

반복문이 중첩되어 있기 때문에 좋은 코드는 아닌 것 같다. 컴퓨터 수가 100 이하로 제한되기 때문에 그리 크지 않다고 판단해서 이렇게 코드를 짰다.

 

 

메모리: 29200 KB

시간: 164 ms

문제

https://www.acmicpc.net/problem/2606

2606번: 바이러스