백준 16940 파이썬 - BFS 스페셜 저지

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

16940번: BFS 스페셜 저지

문제

BOJ에서 정답이 여러가지인 경우에는 스페셜 저지를 사용한다. 스페셜 저지는 유저가 출력한 답을 검증하는 코드를 통해서 정답 유무를 결정하는 방식이다. 오늘은 스페셜 저지 코드를 하나 만들어보려고 한다.

정점의 개수가 N이고, 정점에 1부터 N까지 번호가 매겨져있는 양방향 그래프가 있을 때, BFS 알고리즘은 다음과 같은 형태로 이루어져 있다.

1. 큐에 시작 정점을 넣는다. 이 문제에서 시작 정점은 1이다. 1을 방문했다고 처리한다.
2. 큐가 비어 있지 않은 동안 다음을 반복한다.
   1. 큐에 들어있는 첫 정점을 큐에서 꺼낸다. 이 정점을 x라고 하자.
   2. x와 연결되어 있으면, 아직 방문하지 않은 정점 y를 모두 큐에 넣는다. 모든 y를 방문했다고 처리한다.

2-2 단계에서 방문하지 않은 정점을 방문하는 순서는 중요하지 않다. 따라서, BFS의 결과는 여러가지가 나올 수 있다.

트리가 주어졌을 때, 올바른 BFS 방문 순서인지 구해보자.

코드

import sys
from collections import deque
n = int(sys.stdin.readline())
arr = [set() for _ in range(n + 1)]
arr[0].add(1)
for _ in range(n - 1):
    a, b = map(int, sys.stdin.readline().split())
    arr[a].add(b)
    arr[b].add(a)
def bfs():
    queue = deque()
    queue.append(0)
    idx = 0
    ans = list(map(int, sys.stdin.readline().split()))
    for i in ans:
        while i not in arr[queue[idx]]:
            idx += 1
            if idx == len(queue):
                return 0
        queue.append(i)
    return 1
print(bfs())

풀이

ans에 있는 숫자를 i로 하나씩 가져와 큐 내부의 노드의 arr에 있는지 비교해 없으면 다음 인덱스의 값이 있는지 또 비교하고, 있으면 queue에 그 값을 추가한다. idx값이 queue의 마지막까지 왔는데 i가 arr배열에서 queue와 연결된 다음 숫자에 없다는 것은 BFS로 탐색한 순서가 아니라는 뜻이다.