[백준] 1068 파이썬(python) : 트리

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

1068번: 트리

 

 

시작노드는 인덱스 0번이 아닐 수 있다. 처음엔 무조건 0번 노드가 시작점인줄 알고 풀었다가 틀렸다.

 

처음시도한 틀린 코드)

import sys
sys.setrecursionlimit(10**5)
n = int(sys.stdin.readline())
tree = list(map(int, sys.stdin.readline().split()))
graph = [[] for _ in range(n+1) ]
for idx in range(1, len(tree)):
    graph[tree[idx]].append(idx)
x = int(sys.stdin.readline())
visit = [False]*(n+1)
visit[0] = True
cnt = 0
def dfs(v):
    global cnt
    for i in graph[v]:
        if not visit[i]:
            visit[i] = True
            dfs(i)
    if not graph[v]:
        cnt += 1
for i in range(n):
    if x in graph[i]:
        graph[i].remove(x)
if x == 0:
    print(0)
    exit()
dfs(0)
print(cnt)

 

나의 정답코드)

import sys
sys.setrecursionlimit(10**5)
n = int(sys.stdin.readline())
tree = list(map(int, sys.stdin.readline().split())) #1
graph = [[] for _ in range(n+2) ] #2
for idx in range(n):
    if tree[idx] == -1: #3
        graph[n+1].append(idx)
    else: #4
        graph[tree[idx]].append(idx)
x = int(sys.stdin.readline()) #5
visit = [False]*(n+2)
visit[n+1] = True
for i in range(n+2): #6
    if x in graph[i]:
        graph[i].remove(x)
cnt = 0
def dfs(v):
    global cnt
    if not graph[v]:
        if v == n+1: #7
            return
        cnt += 1 #8
        return
    for i in graph[v]:
        if not visit[i]:
            visit[i] = True
            dfs(i)
dfs(n+1) #9
print(cnt)

 

#1 : 트리의 연결관계를 입력

#2 : 그래프 리스트 선언 (시작노드를 n+1노드에 넣어줄것이기 때문에 n+2로 선언)

#3 : 하나씩 입력값을 비교하는데 입력값이 -1이면 루트노드이기 때문에 n+1번 노드에 현재 인덱스 삽입

#4 : 시작노드가 아니라면 부모노드와 자식노드를 연결

#5 : 삭제할 노드를 입력

#6 : 0부터 n+1번 인덱스까지 반복을 돌면서 내부 리스트에 x가 있으면 삭제

부모노드와의 연결을 끊으면 dfs를 수행할때 끊어진 노드의 자식노드들도 트리와의 연결에서 끊어지게 되므로 부모노드와의 연결만 끊으면 된다.

#7 : dfs수행하는 중 리스트가 비어있고 현재 노드의 리스트가 n+1번 리스트라면 시작노드가 들어있는 리스트이므로 리턴

(예제 입력3과 같이 시작노드를 지워버리면 리프노드는 0이 된다)

#8 : 현재 인덱스의 리스트가 비어있지만 시작노드의 리스트는 아니라면 더 이상 탐색할 수 없는 리프노드이므로 cnt +1 해주고 리턴

#9 : 시작노드는 n+1번 리스트에 들어있기 때문에 dfs(n+1) 실행

 

 

아이디어가 떠오르고 풀릴것 같아서 오랫동안 고민한 문제다. 골드문제를 직접 풀어서 기분이 좋다.