[Python] 1855번 영준이의 진짜 BFS

문제 링크

 

성능 요약

메모리: 115,304 KB, 시간: 9,187 ms, 코드길이: 1,210 Bytes

 

from collections import deque
 
# 큐 순서 정하기
def bfs(x):
    lst=[x]
    queue=deque()
    queue.append(x)
 
    while queue:
        q=queue.popleft()
        for nq in child[q]:
            lst.append(nq)
            queue.append(nq)
 
    return lst        
 
def cal_depth():
    queue=deque()
    queue.append(1)
    d[1]=0
 
    while queue:
        x=queue.popleft()
        for y in child[x]:
            d[y]=d[x]+1
            queue.append(y)
 
def lca(a,b):
    # 깊이 맞추기
    while d[a]!=d[b]:
        if d[a]>d[b]:
            a=parent[a]
        else:
            b=parent[b]
     
    # 공통 조상 찾기
    while a!=b:
        a=parent[a]
        b=parent[b]
 
    return a
 
T=int(input())
for t in range(1,T+1):
    # 1.입력
    N=int(input())
 
    # 2.자료형 정의
    parent=[0,0]+list(map(int,input().split()))
 
    child=[[] for _ in range(N+1)]
    for n in range(1,N+1):
        child[parent[n]].append(n)
 
    d=[0]*(N+1) # 깊이
 
    # 3.동작 알고리즘
    # 1번 노드가 루트
    cal_depth()
    nodes=bfs(1)
 
    ans=0
    for i in range(N-1):
        a=nodes[i]
        b=nodes[i+1]
        sr=lca(a,b) #sub root
        ans+=(d[a]-d[sr])+(d[b]-d[sr])
 
    # 3.출력
    print(f'#{t} {ans}')

 

깊이를 구할 때 DFS를 이용하여 풀었더니 런타임 에러(Runtime Error)가 발생하였다. 이는 DFS 깊이 제한 초과(RecursionError) 때문이다. 그래서 깊이 구하는 방식을 BFS로 바꾸었더니 PASS하였다. 

 

TO DO.

LCA 알고리즘을 이진 리프팅을 통해 개선할 수 있는데 이진 리프팅을 통해 시간과 메모리를 줄인 코드를 보아서, 이진 리프팅으로 LCA 알고리즘을 푸는 연습을 해야겠다.