[ BOJ 11725 ] 트리의 부모 찾기(Python)

문제

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

트리가 주어졌을 때, 각 노드의 부모를 구하는 문제다.
이때 트리의 정점은 1이고, 2번 노드부터 부모를 출력한다.

문제 풀이

입력 받은 정점으로 그래프를 양방향으로 만들어줬다.

for _ in range(n-1):
    x, y = map(int, input().rstrip().rsplit())
    if x not in graph:
        graph[x]=[y]
    else:
        graph[x].append(y)

    if y not in graph:
        graph[y]=[x]
    else:
        graph[y].append(x)

1번 노드와 정답 리스트를 dfs 함수의 인자로 넣어주었다.
answer 리스트의 각 인덱스에는 해당 인덱스의 부모를 넣어줄 것이다.

#dfs(v, answer list)
answer = (dfs(1,answer))[2:]

[ dfs 코드 ]

1. v 노드를 방문 해준다.
2. v와 연결된 u 노드 중 방문 안한 노드가 있다면 방문해준다.
   2-1. 만약 u가 1이 아니라면 dfs의 인자로 넣어준다. u의 부모는 v이다.
3. 만약 u를 방문했는데, u가 1이 아니라면 u의 부모는 v이다.

def dfs(v,answer):
    visited[v]=1
    for u in graph[v]:
        if visited[u]==0:
            if u!=1:
                dfs(u,answer)
                answer[u] = v
        else:
            if u!=1:
                answer[u] = v

    return answer

코드

import sys
sys.setrecursionlimit(100000)
input = sys.stdin.readline

n = int(input().rstrip())
graph = {}
answer = [0] * (n+1)
visited = [0] * (n+1)

for _ in range(n-1):
    x, y = map(int, input().rstrip().rsplit())
    if x not in graph:
        graph[x]=[y]
    else:
        graph[x].append(y)

    if y not in graph:
        graph[y]=[x]
    else:
        graph[y].append(x)

def dfs(v,answer):
    visited[v]=1
    for u in graph[v]:
        if visited[u]==0:
            if u!=1:
                dfs(u,answer)
                answer[u] = v
        else:
            if u!=1:
                answer[u] = v

    return answer
        
answer = (dfs(1,answer))[2:]
for num in answer:
    print(num)