[백준 15681번] 파이썬 - 트리와 쿼리

백준 15681 - 트리와 쿼리

시간 제한 1초, 메모리 제한 128MB

# 조건

• 간선에 가중치와 방향성이 없는 임의의 루트 있는 트리가 주어졌을 때, 아래의 쿼리에 답해보도록 하자.
  • 정점 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줄에 걸쳐 각 쿼리의 답을 정수 하나로 출력한다.

# 접근 방법

• 재귀를 이용하여 풀어주었다.
• 우선 입력받는 간선의 정보를 양방향으로 기록해준다.
• 각 정점의 부모를 기록할 1차원 리스트 parent, 자식을 기록할 2차원 리스트 children, 서브트리에 속하는 정점의 수를 기록할 cnt_sub 리스트를 생성해준다.
• 첫 번째로는, make_tree 함수에 (루트 노드, 부모노드 => -1)을 인자로 전달하여 실행해준다.
  • 처음에 기록한 graph 정보에서 노드들을 순회하며, 부모 노드와 다음 노드가 같지 않다면
  • children List에 기록, parent 리스트 갱신 후에 node=>cur_node, cur_node=>par_node 인자로 전달하여 재귀를 실행한다.
  • 위의 함수가 끝나고 나면 parent 리스트는 각 정점의 부모 노드로 기록되어 있다.
• 두 번째로는 sub_tree 함수를 이용하여 서브 트리에 속한 정점의 수를 구해준다.
  • 루트 노드를 최초로 넣어주며 children 리스트에 기록되어 있는 자식 정점들을 재귀로 들어가준다.
  • leaf 노드에 도착하게 되면 cnt_sub[cur_node] += cnt_sub[node]가 실행되게 되는데
  • 즉, 리프 노드에서부터 해당 자식 노드의 서브트리에 존재하는 정점의 수를 부모노드에 더해주는 것이다.

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

def make_tree(cur_node, par_node):
    for node in graph[cur_node]:
        if node != par_node:
            children[cur_node].append(node)
            parent[node] = cur_node
            make_tree(node, cur_node)

def sub_tree(cur_node):
    for node in children[cur_node]:
        sub_tree(node)
        cnt_sub[cur_node] += cnt_sub[node]


N, R, Q = map(int, input().split())
graph = [[] for _ in range(N+1)]

for _ in range(N-1):
    a, b = map(int, input().split())
    graph[a].append(b)
    graph[b].append(a)

parent = [0] * (N+1)
parent[R] = -1
children = [[] for _ in range(N+1)]
cnt_sub = [1] * (N+1)
make_tree(R, -1)
sub_tree(R)

for _ in range(Q):
    res = int(input())
    print(cnt_sub[res])