Jin의 일상

풀이

import sys
INPUT = sys.stdin.readline


def remove_node(tree, parents, n):
    """ 요청받은 노드를 제거 합니다.

    1. tree 내에 자신(n)을 자식으로 나타내는 부모 노드에서 n을 지웁니다.
      -> parents[n] == -1은 루트 노드이므로 부모 노드가 없습니다.
    2. 재귀적으로 자식 노드를 지웁니다. 지워진 자식 노드의 tree 내에 값을 [-1]로 표기합니다.
    """
    if parents[n] != -1:
        tree[parents[n]].remove(n)

    def remove_child(x):
        for i in tree[x]:
            remove_child(i)
        tree[x] = [-1]
    remove_child(n)


def solution():
    """ 자식이 없는 노드의 개수를 출력합니다.

    1. 이중 리스트(tree)를 생성하고 자신의 자식 노드 인덱스를 넣습니다.
      -> 문제에서 주어진 값 중 -1은 부모 노드가 없음을 의미합니다.
    2. 문제에서 주어진 노드를 제거합니다.
    3. tree에서 자식이 없는 노드([], 빈 노드)의 개수를 출력합니다.
    """
    node_num = int(INPUT())
    parents = list(map(int, INPUT().split()))
    tree = [[] for _ in range(node_num)]

    for i in range(node_num):
        if parents[i] == -1:
            continue
        tree[parents[i]].append(i)

    remove_node(tree, parents, node_to_remove := int(INPUT()))
    return tree.count([])

if __name__ == '__main__':
    print(solution())

반례: ValueError 발생 원인

풀이 과정에서 ValueError가 발생했습니다. 0번 인덱스는 루트 노드(-1)라는 가정했기 때문입니다.

이에 대해 아래와 같은 반례를 고려해볼 수 있습니다.

5

1 2 3 4 -1

4

런타임 발생 당시 코드 일부입니다. n != 0을 확인해주시기 바랍니다. 이를 parents[n] != -1으로 수정했습니다.

import sys
INPUT = sys.stdin.readline


def remove_node(tree, parents, n):
    """ 요청받은 노드를 제거 합니다.

    1. tree 내에 자신(n)을 자식으로 나타내는 부모 노드에서 n을 지웁니다.
      -> n == 0은 루트 노드이므로 부모 노드가 없습니다.
    2. 재귀적으로 자식 노드를 지웁니다. 지워진 자식 노드의 tree 내에 값을 [-1]로 표기합니다.
    """
    if n != 0:
        tree[parents[n]].remove(n)

    def remove_child(x):
        for i in tree[x]:
            remove_child(i)
        tree[x] = [-1]
    remove_child(n)