[Python] 15681 - 트리와 쿼리

[Gold V] 트리와 쿼리 - 15681 

[문제 링크](https://www.acmicpc.net/problem/15681)

🗝️알고리즘 분류

• 트리의 탐색
• 깊이 우선 탐색(DFS)

---

💻문제 정의

트리가 주어졌을 때, 쿼리에 대한 답변을 출력하는 문제이다. 트리는 가중치와 방향성이 없고, 루트는 있다. 쿼리는 정점 U를 루트로 하는 서브트리의 정점의 개수를 구하는 것이다.

 

💡접근 및 설계

루트가 있는 트리를 설계한 후, 쿼리로 부터 오는 정점 U가 루트가 되는 서브트리의 정점을 개수를 구해야 한다.

 

일반적인 dfs로 접근하였고, 각 정점마다 정점이 루트일 때, 서브트리의 개수를 포함시키도록 하였다.

 

✏️알고리즘 풀이

visited = [False for _ in range(N+1)] # 방문 체크 표시 및 서브트리의 개수 저장

def dfs(root):
    visited[root] = 1 # 나 자신 먼저 포함

    for ix in graph[root]: # 자식 노드에 대해서
        if not (visited[ix]): # 방문하지 않았더라면
            visited[root] += dfs(ix) # 나 + 자식 노드의 서브트리 정점 개수

    return visited[root] # 정점 트리 개수 리턴

dfs(R) # 루트가 R일 때 모든 정점에 대해서 서브트리의 정점의 개수가 visited에 저장

graph는 양방향 그래프로 먼저 설계하였다. 입력에서 방향은 주어지지 않고 두 정점에 대한 간선 정보만 주어지기 때문이다.

 

R은 graph의 루트 노드이고, 루트로 부터 leaf 노드 까지 dfs를 한다. leaf 노드에서는 더 이상 dfs를 실행할 수 없으므로 visited[leaf] = 1을 반환한다. 이후부터 부모 노드들은 자식 노드 + 자기 자신을 하여 자기 자신에 저장된다. 최종적으로 루트 노드에는 자기 자신을 포함한 모든 정점의 개수가 visited에 저장된다.

 

쿼리에 해당하는 값을 visited로 호출하면 문제 해결.

 

🗒️전체 코드

# 15681 트리와 쿼리

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

N, R, Q = map(int, input().split())

# 양방향 그래프 생성
graph = [[] for _ in range(N+1)]
for _ in range(N-1):
    u, v = map(int, input().split())
    graph[u].append(v)
    graph[v].append(u)

visited = [False for _ in range(N+1)] # 방문 체크 표시 및 서브트리의 개수 저장

def dfs(root):
    visited[root] = 1 # 나 자신 먼저 포함

    for ix in graph[root]: # 자식 노드에 대해서
        if not (visited[ix]): # 방문하지 않았더라면
            visited[root] += dfs(ix) # 나 + 자식 노드의 서브트리 정점 개수

    return visited[root] # 정점 트리 개수 리턴

dfs(R) # 루트가 R일 때 모든 정점에 대해서 서브트리의 정점의 개수가 visited에 저장

for _ in range(Q):
    print(visited[int(input())])

---