16964번 DFS 스페셜 저지 백준 파이썬

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

16964번: DFS 스페셜 저지

 

이 문제를 풀었을 때 75%에서 도저히 개선이 안되서 결국 다른 분의 코드를 참고한 문제였다.

 

 

 

 

import sys
N = int(sys.stdin.readline())
graph=[[] for _ in range(N+1)]
for _ in range(N-1):
    a,b=map(int,sys.stdin.readline().split())
    graph[a].append(b)
    graph[b].append(a)
ans= list(map(int,sys.stdin.readline().split()))
for i in range(1,N+1):
    graph[i].sort()
visited = [False]*(N+1)
visited[0]=True
go=1
 
flag=False
def dfs(n):
    global go,flag
    stk=[]
    visited[n]=True
    if go==N:
        if all(visited):
            flag=True
            print(1)
            sys.exit(0)
            return True
        else:
            return False
 
    for i in graph[n]:
        if not visited[i]:
            stk.append(i)
    while stk:
        if ans[go] in stk:
            visited[ans[go]]=True
            stk.remove(ans[go])
            go+=1
            dfs(ans[go-1])
        else:
            print(0)
            sys.exit(0)
            return False
            
   
       
if ans[0]==1:
    dfs(1)
else:
    print(0)

위의 코드는 시간초과로 인하여 틀린 코드이다.

최악의 경우를 생각해보자. 예를 들어 1의 주변으로 2부터 99999까지의 수가 다 연결되어 있는 그래프라 생각해보자.

그리고 순서가 99999 부터 2순으로 되어있다고 생각해보자. 근데 graph는 정렬이 되어있기에 2부터 99999순으로 저장되어 있다. 그러면 만약 99999까지 찾는데 대략 십만, 그리고 99998까지도 대략 십만 이러면 안봐도 시간초과가 날 수 밖에 없는 구조다.

그래서 밑의 코드는 그것을 극복하기 위해 각각의 깊이에 따른 레벨을 두고 자신을 포함한 총 길이를 각 리스트에 저장해놓고 그 결과값을 이용해 답을 도출해내는 코드이다.

 

import sys
N = int(sys.stdin.readline())
graph=[[] for _ in range(N+1)]
for _ in range(N-1):
    a,b=map(int,sys.stdin.readline().split())
    graph[a].append(b)
    graph[b].append(a)
ans= list(map(int,sys.stdin.readline().split()))
level=[False]*(N+1)
tsize = [0]*(N+1)
visited = [False]*(N+1)

def dfs(x,lv):
    if visited[x]:
        return 0
    visited[x]=True
    size=1
    level[x]=lv
    for i in range(len(graph[x])):
        next =graph[x][i]
        size+=dfs(next,lv+1)
    tsize[x]=size
    return size


if ans[0]!=1:
    print("0")
    sys.exit(0)
else:
    dfs(1,0)
    for i in range(1,N):
        x=ans[i]
        if tsize[x] == 1 or i + tsize[x] >=N:
            continue
        next = ans[i+tsize[x]]
        if level[next]>level[x]:
            print(0)
            sys.exit(0)
    print(1)

참고:https://blog.joonas.io/94

BOJ 16964 - DFS 스페셜 저지

 

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