[백준] 1068번 트리 파이썬

 

문제

1068번: 트리

풀이

1. 노드 클래스와 트리 클래스를 기본적으로 틀을 짜줍니다.

2. 트리 클래스에 깊이 우선 탐색 알고리즘을 이용하여 현재 노드에서 자식노드가 지워줄 노드가 같다면 해당 자식을 리스트에서 제거합니다. 

3. 이때 루트 노드와 지워줄 노드가 같다면 트리의 root를 none으로 바꿔주고 예외 처리로 0을 출력해줍니다. ( 루트 노드부터 지운다면 트리 전체가 없어진다는 뜻)

4. 그후 다시 트리 전체를 깊이 우선 탐색하여 현재의 노드가 자식이 없을 경우에 결과값을 카운팅해줍니다.

5. 탐색 종료후 카운팅된 결과값을 출력합니다.

 

** 루트 노드가 단말노드가 될 경우를 처리해 주어여 완벽히 통과가 가능합니다! **

 

코드

class Node:
    def __init__(self, data):
        self.data = data
        self.childs = list()


class Tree:
    def __init__(self, n):
        self.root = None
        self.nodes = list(Node(i) for i in range(n))

    def make_tree(self, parent, child):
        parent.childs.append(child)

    def cut(self, cut_node):
        stack = [tree.root]
        if tree.root == cut_node:
            tree.root = None
            return
        while stack:
            for _ in range(len(stack)):
                node = stack.pop()
                for child in node.childs:
                    if child == cut_node:
                        node.childs.remove(cut_node)
                        return
                    stack.append(child)

    def find_leaves(self):
        leaves = 0
        stack = [tree.root]
        while stack:
            for _ in range(len(stack)):
                node = stack.pop()
                if node.childs == []:
                    leaves += 1
                    continue
                for child in node.childs:
                    stack.append(child)
        return leaves


n = int(input())
tree = Tree(n)

for child, parent in enumerate(list(map(int, input().split()))):
    if parent == -1:
        tree.root = tree.nodes[child]
    else:
        tree.make_tree(tree.nodes[parent], tree.nodes[child])
        
cut_node = tree.nodes[int(input())]
tree.cut(cut_node)

if tree.root is None:
    print(0)
else:
    print(tree.find_leaves())