Baekjoon 15681. 트리와 쿼리 / Python

15681. 트리와 쿼리

난이도 : 골드 5
소요 시간 : 10분
날짜 : 2024.12.23
언어 : 파이썬
알고리즘 유형 : dfs, dp 트리, 그래프

설명 보기전에 문제 풀어보러 가기

1. 문제 설명

1. 트리 구조의 그래프가 주어지고, 루트 노드가 지정된다.
2. 각 노드 u를 루트로 하는 서브트리의 정점 수를 계산한다.

2. 해결 방식

1. 자기 자신을 포함
2. 자식노드를 탐색하고 결과를 누적시킨다.

3. 정답 코드

import sys
sys.stdin = open("C:/Users/ghtjd/Desktop/tmp/python/input.txt", "r")

sys.setrecursionlimit(1e10)
# input = sys.stdin.readline
n, r, q = map(int, input().split()) # 2 ≤ N ≤ 10^5, 1 ≤ R ≤ N, 1 ≤ Q ≤ 10^5

res = [0] * (n + 1)
tree = [[] for _ in range(n + 1)]

# 루트있는 트리
for _ in range(n - 1):
    u, v = map(int, input().split())
    tree[u].append(v)
    tree[v].append(u)

# 서브트리 만들기
def dfs(x):
    res[x] = 1
    for i in tree[x]:
        if not res[i]:
            dfs(i)
            res[x] += res[i]
dfs(r)

# u를 루트로 하는 서브트리의 정점 수 출력
for _ in range(q):
    u = int(input())
    print(res[u])