[백준 1068번] 트리 (python)

문제

트리에서 리프 노드란, 자식의 개수가 0인 노드를 말한다.

트리가 주어졌을 때, 노드 하나를 지울 것이다. 그 때, 남은 트리에서 리프 노드의 개수를 구하는 프로그램을 작성하시오. 노드를 지우면 그 노드와 노드의 모든 자손이 트리에서 제거된다.

예를 들어, 다음과 같은 트리가 있다고 하자.

현재 리프 노드의 개수는 3개이다. (초록색 색칠된 노드) 이때, 1번을 지우면, 다음과 같이 변한다. 검정색으로 색칠된 노드가 트리에서 제거된 노드이다.

이제 리프 노드의 개수는 1개이다.

 

기록

(2일차 트리) 틀렸습니다 지옥............. 생각을 좀 넓게 하자!

 

코드

<잘못된코드>

from collections import defaultdict

n = int(input())
note = list(map(int, input().split()))
delete_node = int(input())

if delete_node==0: print(0)
else :
    tree = {0:[]}
    for i in range(1, len(note)):
        if i==delete_node: continue
        if note[i] in tree:
            tree[i] = []
            tree[note[i]].append(i)

    leaf = 0
    for root, nodes in tree.items():
        if len(nodes)==0 : leaf+=1
    print(leaf)

문제 : 0이 무조건 루트라고 생각했고, 노드가 트리 순서대로 입력된다고 생각했던 것..

해결 : 각 노드의 연결정보를 dict로 만들어 놓고 그 뒤에 leaf 찾기 하도록 변경

 

<고친코드>

from collections import defaultdict

# 변수 초기화
n = int(input())
note = list(map(int, input().split()))
delete_node = int(input())
tree = defaultdict(list)
for i in range(len(note)):
    tree[note[i]].append(i)

# delete_node 제외한 트리 leaf 찾기
def count_leaf(tree, node, delete_node):

    if len(tree[node])==0 : return 1
    if len(tree[node])==1 and tree[node][0]==delete_node and node!=-1 : return 1

    count = 0
    for n in tree[node]:
        if n != delete_node:
            count += count_leaf(tree, n, delete_node)
    return count

print(count_leaf(tree, -1, delete_node))

각 노드 연결정보 미리 만들어놓기 -> dfs로 delete node 제외한 경로로만 탐색해서 leaf 찾기 해결!!