2606 바이러스 (python)

문제

 

 

첫 번째 풀이

import sys

input = sys.stdin.readline

computer = int(input().strip())

n = int(input().strip())

connect = []

for _ in range(n):

    connect.append(list(map(int, input().strip().split())))

virus = [0] * (computer + 1)

virus[1] = 1

def spread_virus(start):

    queue = [start]

   

    while queue:

        current = queue.pop(0)

        for a, b in connect:

            if a == current and virus[b] == 0:

                virus[b] = 1

                queue.append(b)

            elif b == current and virus[a] == 0:

                virus[a] = 1

                queue.append(a)

spread_virus(1)

cnt = sum(virus) - 1

print(cnt)

 

어제 친구가 BFS문제를 풀 때 왜 deque를 사용하냐고 물어보았다.

deque를 쓰는 이유는 pop(0)와 같은 메서드를 수행할 때

리스트의 경우 O(N)연산을 수행하지만 deque는 O(1) 연산을 수행하기 때문에

리스트보다 deque를 사용한다고 말해주었다.

 

갑자기 deque를 사용하지 않고 문제를 풀어보고 싶어서 첫 번째 풀이는 사용하지 않고 풀었다.

 

코드 해설

virus 배열은 감염 된 컴퓨터를 체크하는 배열

0이면 감염X 1이면 감염, 0으로 초기화

 

문제에서 1번 컴퓨터가 웜 바이러스에 걸렸을 때 전염 컴퓨터를 찾는 것이므로

첫번째 virus[1] = 1

 

def spread_virus 함수로 BFS탐색을 해  virus에 걸리면 virus 배열에서

바이러스가 걸린 컴퓨터들(인덱스)의 값을 1로 바꿔주었다.

 

cnt = sum(virus) - 1 # 1번 컴퓨터 제외

print(cnt)로 cnt를 출력해주었다.

 

 

두 번째 풀이

import sys

from collections import deque

input = sys.stdin.readline

computer = int(input().strip())

n = int(input().strip())

connect = []

for _ in range(n):

    connect.append(list(map(int, input().strip().split())))

virus = [0] * (computer + 1)

virus[1] = 1

def spread_virus(start):

    queue = deque([start])

   

    while queue:

        current = queue.popleft()

        for a, b in connect:

            if a == current and virus[b] == 0:

                virus[b] = 1

                queue.append(b)

            elif b == current and virus[a] == 0:

                virus[a] = 1

                queue.append(a)

spread_virus(1)

cnt = sum(virus) - 1

print(cnt)

 

그냥 deque만 추가했다.

 

제출

 

근데 deque를 사용하지 않은 풀이가 메모리와 시간복잡도가 적었다.

알아보니 입력 크기가 작을 경우, 리스트의 pop()연산의 O(n) 시간복잡도가 크게 문제가 되지 않고,

오히려 deque의 추가적인 초기화 및 메서드 호출 비용이 오히려 메모리나 시간 복잡도가 더 걸릴 수 있게 만든다는 것을 알게 되었다.