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

문제:www.acmicpc.net/problem/16940

16940번: BFS 스페셜 저지

 

import sys
from collections import deque
N = int(sys.stdin.readline().rstrip())
graph=[[] for _ in range(N+1)]
for _ in range(N-1):
    a,b= map(int,sys.stdin.readline().rstrip().split())
    graph[a].append(b)
    graph[b].append(a)
for i in range(1,N+1):
    graph[i].sort()
ans=list(map(int,sys.stdin.readline().split()))
stk=[]
flag=False
visited = [False]*(N+1)
def bfs(n):
    global flag
    q =deque()
    visited[n]=True
    q.append(n)
    start=1
    tmp=[]
    while q:
        stk=[]
        for i in graph[q.popleft()]:
            if not visited[i]:
                visited[i]=True
                stk.append(i)
        li=ans[start:start+len(stk)]
        tmp=list(li)
        start=start+len(stk)
        li.sort()
        if li!=stk:
            flag=True
            break
        else:
            for i in tmp:
                q.append(i)
if ans[0]==1:       
    bfs(1)
    print(1) if not flag else print(0)
else:
    print(0)
 

10
1 2
2 3
2 4
3 5
1 7
7 10
6 8
6 9
1 6
1 7 6 2 10 9 8 4 3 5

위와 같은 그림의 로직대로 흘러간다.

우선 같은 층에 있는 리스트를 부분적으로 잘라내고 비교하고자 하는 리스트 부분을 정렬해서 조건문에 확인하고 만일 다르다면 flag=True break

 

그리고 마지막에 ans[0]이 1이 아닐 경우에는 무조건 0을 출력하도록 한다.