백준 15681번: 트리와 쿼리 (Python)

힌트

접근

문제에 주어진 힌트에 문제를 풀기 위한 간략한 트리에 대한 설명이 나온다. 어찌보면 간단한 내용이지만 이제 배워가는 나에게는 큰 도움이 된 것 같다.

문제 풀이에 도움이 되는 힌트의 핵심은 트리의 정점이 가진 서브트리의 갯수를 저장해주라는 것이다. 루트를 시작으로 하여 서브트리가 계속 아래로 연결되어 확장되는 식이기 때문에 BFS보다는 DFS 알고리즘을 이용하는 것이 직관적으로 풀 수 있는 방법이었던 것 같다.

DFS 함수를 만들고 루트를 시작으로 루트에서 이어진 자식 노드들을 탐색하여 자식 노드들이 가진 서브트리 정점의 갯수를 더해서 저장해준다. 그 과정에서 자식 노드들이 가진 서브트리의 갯수가 필요한데, 이 것을 재귀적으로 반복하여 결과적으로는 가장 깊은 곳에 있는 노드들의 서브트리부터 채워주면서 마지막으로 루트 노드의 서브트리 갯수가 채워지게 된다.

마지막으로 입력 받는 쿼리들에 해당하는 서브트리 갯수를 출력만 해주면 되는 것이다.

코드

import sys

sys.setrecursionlimit(10 ** 9)
n, r, q = map(int, sys.stdin.readline().split())
lines = [[] for _ in range(n + 1)]
subtree = [0] * (n + 1)  # i 번 노드를 루트로 하는 서브트리의 정점의 갯수
visited = [False] * (n + 1)
for _ in range(n - 1):
    a, b = map(int, sys.stdin.readline().split())
    lines[a].append(b)
    lines[b].append(a)


def dfs(root):
    subtree[root] = 1
    visited[root] = True
    for i in lines[root]:
        if not visited[i]:
            dfs(i)
            subtree[root] += subtree[i]
    return


dfs(r)

for _ in range(q):
    print(subtree[int(sys.stdin.readline())])

더 생각해 볼 것?

...

코드나 내용 관련 조언, 부족한 점 및 질문 언제든 환영합니다!