Koala - 14기/코딩테스트 준비 스터디

[백준/python] 2644 촌수계산

ㄱㅈㅅㅇ 2024. 5. 5. 18:07

2644번: 촌수계산 (acmicpc.net)

 

코드 

import sys
input = sys.stdin.readline
n=int(input())
a,b=map(int,input().split())
m=int(input())
di={}
for _ in range(m): #di 입력받기
    u,v=map(int,input().split())
    if u in di: di[u].append(v)
    else: di[u]=[v]
    if v in di: di[v].append(u)
    else: di[v]=[u]

ans_visit={}
st=[a]

while st:
    node=st.pop()
    if node in ans_visit: continue
    st.extend(di[node])
    if node==a: ans_visit[node] = a
    else:
        for x in di[node]:
            if x in ans_visit:
                ans_visit[node] = x

if b not in ans_visit: print(-1)
else:
    ans = 0
    while 1:
        b=ans_visit[b]
        ans+=1
        if b==a:
            break
    print(ans)

풀이

모든 관계를 di에 입력받고, dfs 또는 bfs를 돌리면 된다. 이 과정에서 visit에 부모 관계임을 알도록 ans_visit[node]=x해준다. 우리가 알고싶은 관계는 a, b이니까 a부터 시작한다. a부터 시작하여 모두 방문했는데, visit에 b가 없다면 둘은 친척이 아니다. -1출력. 그게 아니라면 visit에 저장되어있는걸 본다. b부터 시작해서 부모 부모에게로 올라가고 그러면서 촌수를 +1해준다. 마침내 a를 만나면 멈추고 촌수를 출력한다.