백준 11725-트리의 부모 찾기 (파이썬)

문제

루트 없는 트리가 주어진다. 이때, 트리의 루트를 1이라고 정했을 때, 각 노드의 부모를 구하는 프로그램을 작성하시오.

입력

첫째 줄에 노드의 개수 N (2 ≤ N ≤ 100,000)이 주어진다. 둘째 줄부터 N-1개의 줄에 트리 상에서 연결된 두 정점이 주어진다.

출력

첫째 줄부터 N-1개의 줄에 각 노드의 부모 노드 번호를 2번 노드부터 순서대로 출력한다.

예제 입력 1

7
1 6
6 3
3 5
4 1
2 4
4 7

예제 출력 1

4
6
1
3
1
4

예제 입력 2

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

예제 출력 2 

1
1
2
3
3
4
4
5
5
6
6

---

오답 코드1

import sys

input = sys.stdin.readline
n = int(input())

# 인접리스트 생성
adj = [[] for _ in range(n+1)]
for _ in range(n - 1):
    n1, n2 = map(int, input().split())
    adj[n1].append(n2)
    adj[n2].append(n1)
for i in range(2, n+1):
    print(adj[i][0])

첫번째 떠올린 아이디어는 인접리스트를 만들어 각 정점의 첫번째 요소를 부모로 간주해 출력하는 것이다.

주어진 테스트 케이스는 전부 통과했지만 제출하니 오답처리가 되었다.  문제는 입력이 1번 정점부터 나오면 문제 없지만 다른게 먼저 나오면 오답이 나오는 것이었다. 예시는 모두 1번 정점부터 주어졌기 때문에 통과할 수 있었던 것..

오답 코드2(시간 초과)

import sys

sys.setrecursionlimit(10 ** 9)
input = sys.stdin.readline
n = int(input())

# 인접리스트 생성
adj = [[] for _ in range(n + 1)]
for _ in range(n - 1):
    n1, n2 = map(int, input().split())
    adj[n1].append(n2)
    adj[n2].append(n1)


def dfs(now, target):
    for nxt in adj[now]:
        if nxt == target:
            print(now)
            return
        if not chk[nxt]:
            chk[nxt] = True
            dfs(nxt, target)


for i in range(2, n + 1):
    chk = [False] * (n + 1)
    dfs(now=1, target=i)

 

그 다음 아이디어는 DFS를 이용해 탐색하는 방법이다. 루트 노드는 1이기 때문에 1부터 DFS로 탐색하면 찾고자하는 모든 정점을 찾을 수 있다. 하지만 DFS함수 매개변수로 목표노드를 전달하고 그 노드를 발견했을 때 다음 목표노드를 찾기위해 N번 만큼 반복하기 때문에 시간복잡도는 O(N^2)로 최악의 경우 100억에에 가까운 연산을 해야하기에 시간초과가 발생한 것이다.

 

정답 코드

import sys

sys.setrecursionlimit(10 ** 9)
input = sys.stdin.readline
n = int(input())

# 인접리스트 생성
adj = [[] for _ in range(n + 1)]
for _ in range(n - 1):
    n1, n2 = map(int, input().split())
    adj[n1].append(n2)
    adj[n2].append(n1)

chk = [0] * (n + 1)


def dfs(now):
    for nxt in adj[now]:
        if not chk[nxt]:
            chk[nxt] = now
            dfs(nxt)


dfs(1)

for i in range(2, n + 1):
    print(chk[i])

 

 

dfs는 모든 노드를 방문하는 완전탐색이기 때문에 이 점을 이용해서 dfs를 한번만 호출하면 되지 않을까라는 생각이 들었다.

이번 문제의 그래프는 1번 노드를 루트로 하는 트리 구조이기 때문에 연결된 모든 노드는 부모-자식관계를 이룬다. 다음 노드로 이동하기 전에 chk[방문할 노드]의 값을 True로 바꾸게 되는데 이 때 True가 아닌 현재 노드 값을 넣어주면 dfs를 한번만 호출하고도 모든 요소의 부모 노드를 기록할 수 있다.

 

 

 

피드백

오답을 고치는 흐름은 잘잡았던 것 같다. 다른 풀이도 찾아보니 인접행렬로 구현하면 메모리 초과가 나온다고 하는데 파이썬 리스트의 사이즈를 알아보고 직접 구현도 해보자.

그리고 양방향 간선으로 구현했던 걸 단방향으로 구현해보고 추가로 bfs로도 풀어본 후에 포스팅하자