[백준 / 1068] 트리 - Python

<접근>

입력받은 트리의 구조를 그린다. 부모 -> 자식으로의 방향이 있는 그래프이다.

지워야 하는 노드를 입력받고, 해당 노드에 연결된 자식 노드를 삭제한다. 만약 해당 자식 노드에 연결된 노드가 없으면 그냥 지우고, 자식 노드가 있을 경우 같은 방법을 반복한다.

 

아래와 같은 예제에서 틀릴 수 있으니 주의

2
-1 0
1

<입력>

N = int(input())
node = list(map(int, input().split()))

tree = {i: [] for i in range(N)}

 

총노드의 개수 N과 각 노드의 부모가 입력되는 node를 입력받는다. 만약 '-1 0 0 1 1'이 입력된다면 0 노드는 최상위 노드, 1, 2번 노드는 0번 노드의 자식이고, 3, 4번 노드는 1번 노드의 자식인 것이다.

 

tree는 해당 노드들의 그래프를 그리기 위해서 딕셔너리 형태로 선언해주고, 노드의 개수만큼 key를 만들어 준다.

<그래프 그리기>

for i in range(N):
    parent = node[i]
    if parent != -1:
        tree[parent].append(i)

delete_node = int(input())

del_leaf(tree, node, delete_node)

 

각 노드의 부모가 담겨있는 node를 조회하면서 최상위 노드(-1)가 아니면 해당 부모에 자식을 추가해준다. 즉, 1번 노드의 부모가 0이면, 노드 0에 1 노드를 tree 그래프에 추가해준다.

즉 '-1 0 0 1 1'이 입력된다면 '{0: [1, 2], 1: [3, 4], 2: [ ], 3: [ ], 4: [ ]}' 트리가 된다.

 

삭제할 노드를 delete_node에 입력받고 함수를 호출해서 함수 del_leaf를 호출한다.

<함수1>

def del_leaf(tree, node, delete_node):
    if len(tree[delete_node]) == 0:
        parent = node[delete_node]
        if parent != -1:
            tree[parent].remove(delete_node)

        del tree[delete_node]

 

만약 입력받은 삭제노드에 자식이 없다면(len(tree [delete_node])==0) 해당 노드를 그래프에서 삭제하고, 해당 노드와 연결된 부모 노드의 값에서 해당 노드의 값(value)을 삭제해야 한다. 먼저 부모 노드가 무엇인지 입력받는 node 리스트에서 찾아내고 tree그래프에서 부모 노드의 key에서 삭제할 노드의 값을 지워준다. 그리고 tree그래프에서 해당 노드의 key를 삭제해준다.

<함수2>

    else:
        leaf_list = tree[delete_node][:]
        for leaf in leaf_list:
            del_leaf(tree, node, leaf)

        parent = node[delete_node]
        if parent != -1:
            tree[parent].remove(delete_node)
        del tree[delete_node]

 

만약 자식노드가 있는 경우 자식노드가 무엇인지 leaf_list에 복사한다. 자식리스트를 돌면서 해당 작업을 똑같이 반복해주기 위함이다. 자식노드도 그래프에서 삭제해줘야하므로 해당 작업을 반복해준다. 자식 노드도 다 삭제하면 해당 노드의 값도 부모 tree의 key에서 제거해주고, tree그래프에서 삭제 노드의 key도 제거해준다.

<출력>

count = 0

for leaf in tree:
    if len(tree[leaf]) == 0:
        count += 1
        
print(count)

리프 노드(leaf node)는 연결된 자식이 없는 노드이므로, tree그래프의 key를 탐색하면서 리스트의 크기가 0인 노드를 발견하면 count를 세준다.

 

<전체코드>

def del_leaf(tree, node, delete_node):
    if len(tree[delete_node]) == 0:
        parent = node[delete_node]
        if parent != -1:
            tree[parent].remove(delete_node)

        del tree[delete_node]
    else:
        leaf_list = tree[delete_node][:]
        for leaf in leaf_list:
            del_leaf(tree, node, leaf)

        parent = node[delete_node]
        if parent != -1:
            tree[parent].remove(delete_node)
        del tree[delete_node]


N = int(input())
node = list(map(int, input().split()))

tree = {i: [] for i in range(N)}

for i in range(N):
    parent = node[i]
    if parent != -1:
        tree[parent].append(i)

delete_node = int(input())

del_leaf(tree, node, delete_node)

count = 0

for leaf in tree:
    if len(tree[leaf]) == 0:
        count += 1
        
print(count)