For a better world

문제 링크

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

풀이 코드

import sys
input = sys.stdin.readline

def dfs(K: int) -> None:
    tree[K] = -2
    for i in range(N):
        if K == tree[i]:
            dfs(i)

if __name__ == "__main__":
    # input 
    N = int(input())
    tree = list(map(int, input().split()))
    K = int(input())

    # main algorithm
    dfs(K)

    # output
    count = 0 
    for i in range(len(tree)):
        if tree[i] != -2 and i not in tree: # 제거 표시인 -2도 아니고, 자신을 부모 노드로 갖는 노드도 존재하지 않을 경우
            count +=1
    print(count)

풀이 과정

핵심 아이디어는 DFS 알고리즘을 통해 부모 노드를 제거 한 뒤, 부모 노드에 종속된 자식 노드를 탐색하며 제거해주는 것이다.

가령 예제 입력 4번으로 예시를 들면 다음과 같다.

[-1, 0, 0, 2, 2, 4, 4, 6, 6]은 0~8까지의 노드 번호 위치에 부모 노드를 나타낸 것이다. 이 중 4번 노드를 없애고자 한다면 5번째 인덱스 값인 2를 제거해주어야 한다. 하지만 실제로 리스트에서 제거할 경우 인덱스의 변화가 생겨 일관된 알고리즘 적용에 번거로워지므로 제거했단 표기로 대신한다. -1은 루트 노드로 사용되니 -2란 값이 적당할 것이고 이를 적용하면 다음과 같아진다. [-1, 0, 0, 2, -2, 4, 4, 6, 6]

위 그림에서 노드 4를 -2로 제거 표기 해준 것이다. 이후 4에 종속된 5, 6 그리고 6에 또 종속된 7, 8을 DFS 알고리즘으로 반복 탐색하며 -2로 변경한다면 4에 종속되지 않았던 노드들만 남아 있게 될 것이다.