백준 부모자식이 있는 트리 탐색

서론

트리와 쿼리

기본적인 트리 탐색에 대한 문제이고 어렵지 않았지만, 입력받은 후에 루트 노드를 주었을 때 다시 트리를 만드는 게 처음이었고, 함수를 재귀하는 과정에서 뭔가 제가 생각하지 못한 방법이 나와서 기록해봅니다.

제한 조건

• 시간제한 1초

시간제한이 1초이고 정점의 수가 10^5까지이기 때문에 쿼리를 받을 때마다 탐색을 하면 안되고 DP를 사용하여 미리 sub tree의 크기를 저장해 놓아야 합니다. 문제 카테고리 자체가 DP이어서 쉽게 떠올렸지만 이런 주어진 조건을 보고 어떻게 풀어야 할지 떠올리는 게 능력이겠네요.

• 메모리 제한 128MB

인풋을 저장해놓고 트리를 만들고 DP를 만들고 하는데 메모리 제한은 문제없습니다. 하지만 sub tree를 count 하는데 recursion을 사용하는데 이 부분에서만 메모리 에러가 나지 않게 하면 될 것 같습니다. 근데 지금 확인해보니 BOJ에서는 Recursion Error를 따로 주네요. 메모리는 문제없습니다.

트리탐색 DFS

다른 글들을 보니까 굳이 tree를 만들지 않더라도 visited 배열로 부모노드를 탐색하지 않게 하는 방법도 있네요. 근데 그냥 입력받아놓고 child만 가지는 tree를 만드는 게 깔끔한 것 같습니다.

import sys
sys.setrecursionlimit(10**6)
input = sys.stdin.readline


def makeTree(currentNode, parent):
    for node in connect[currentNode]:
        if node != parent:
            tree[currentNode].append(node)
            makeTree(node, currentNode)


def countSubtreeNodes(currentNode):
    for node in tree[currentNode]:
        countSubtreeNodes(node)
        dp[currentNode] += dp[node]


# input1
N, R, Q = map(int, input().split())
connect = {}
for i in range(1, N+1):
    connect[i] = []
for __ in range(N-1):
    U, V = map(int, input().split())
    connect[U].append(V)
    connect[V].append(U)

# 받은 인풋으로 tree랑 dp만듬
tree = {}
for i in range(1, N+1):
    tree[i] = []
dp = [1] * (N+1)
makeTree(R, -1)
countSubtreeNodes(R)

# input2
for __ in range(Q):
    print(dp[int(input())])

생각

def countSubtreeNodes(currentNode):
    for node in tree[currentNode]:
        # 먼저 재귀를 부르고
        countSubtreeNodes(node)
        # 그다음에 재귀로 나온 값을 사용한다.
        dp[currentNode] += dp[node]

저는 이 부분이 조금 신기했습니다. 별거 아닐 수 있지만 제가 recursive로 코딩을 많이 안 해봐서 그런지 신기했습니다.
그러니까 재귀함수에서 재귀한 후에 뭔가를 수행하는 코드를 적으면 먼저 재귀를 하러 들어가고 그다음 재귀를 끝내면서 값을 리턴하면서 후에 일어나는 일들의 값을 먼저 사용할 수 있다는 점이 조금 신기했습니다. 말로 적으니까 설명이 힘들지만 재활용할 수 있을 거 같습니다.