[백준 15681 🥇] 트리와 쿼리

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

15681번: 트리와 쿼리

문제

간선에 가중치와 방향성이 없는 임의의 루트 있는 트리가 주어졌을 때, 아래의 쿼리에 답해보도록 하자.

• 정점 U를 루트로 하는 서브트리에 속한 정점의 수를 출력한다.

만약 이 문제를 해결하는 데에 어려움이 있다면, 하단의 힌트에 첨부한 문서를 참고하자.

입력

트리의 정점의 수 N과 루트의 번호 R, 쿼리의 수 Q가 주어진다. (2 ≤ N ≤ 105, 1 ≤ R ≤ N, 1 ≤ Q ≤ 105)

이어 N-1줄에 걸쳐, U V의 형태로 트리에 속한 간선의 정보가 주어진다. (1 ≤ U, V ≤ N, U ≠ V)

이는 U와 V를 양 끝점으로 하는 간선이 트리에 속함을 의미한다.

이어 Q줄에 걸쳐, 문제에 설명한 U가 하나씩 주어진다. (1 ≤ U ≤ N)

입력으로 주어지는 트리는 항상 올바른 트리임이 보장된다.

출력

Q줄에 걸쳐 각 쿼리의 답을 정수 하나로 출력한다.

 

문제풀이

import sys
sys.setrecursionlimit(100000)

n,r,q=map(int, sys.stdin.readline().split(' '))
graph=[[] for _ in range(n+1)]
for _ in range(n-1):
    a,b=map(int, sys.stdin.readline().split(' '))
    graph[a].append(b)
    graph[b].append(a)

visited=[0]*(n+1)
def dfs(v):
    visited[v]+=1
    for node in graph[v]:
        if visited[node]==0:
            visited[v]+=dfs(node)
    return visited[v]

dfs(r)

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

오랜만에 DFS를 이용한 문제풀이였다.

dfs 함수를 만들어서 재귀 함수로 탐색하면서, 각 노드에 대해서 연결된 자식의 수를 visited 배열에 추가해주는 방식으로 구현했다. 

이러한 과정을 모두 마치고 나면, visited 배열에 각 노드별로 자신을 포함한 자식의 수가 저장된다.