[백준/Python] 15681. 트리와 쿼리

15681. 트리와 쿼리

15681번: 트리와 쿼리

• 트리
• DFS
• DP

힌트에서 제공하는 트리의 특성

"트리는 단 하나의 사이클도 존재하지 않는 그래프"

따라서 루트가 정해지면, 특정 노드에 도달하는 최단 경로는 한 가지 뿐이다.

"한 노드와만 연결되어 있는 모든 노드가 루트가 될 수 있다"

또한 루트가 정해지면 노드 간에 부모 자식 관계를 정의할 수 있다.

"아무 정점이든 root로 잡아 subtree를 구성할 수 있다"

따라서 tree는 재귀를 이용하여 구성할 수 있다.

 

---

해당 문제에서 사용하는 두 가지 핵심 알고리즘

그래프 간 연결 관계, 루트를 통해 트리를 구성 (DFS)

여기서 트리를 구성한다는 것은, 양방향 그래프를 부모→자식의 방향만 담는 단방향 그래프로 재구성함을 말한다.

이를 통해 각 노드를 루트로 하는 트리를 쉽게 얻을 수 있다.

def make_tree(graph, current, parent, tree):
        for i in graph[current]:
                if i != parent:
                        tree[current].append(i)
                        make_tree(graph, i, current, tree)

# n행 n열 또는 n+1행 n+1열의 이차원 리스트 graph, tree
make_tree(graph, root, -1, tree)

모든 노드를 루트로 하는 서브트리의 정점 수 구하기 (Top-Down DP)

def count_node(tree, current, size):
        for i in tree[current]:
                count_node(tree, i, size) # child node의 size를 setting
                size[current] += size[i]

size = [1] * (node+1) # 모두 자기 자신을 포함하므로 1로 초기화
count_node(tree, root, size)

전체 코드

import sys
input = sys.stdin.readline
sys.setrecursionlimit(int(1e9))


def make_tree(graph, current, parent, tree):
    for i in graph[current]:
        if i != parent:
            tree[current].append(i)
            make_tree(graph, i, current, tree)


def count_node(tree, current, size):
    for i in tree[current]:
        count_node(tree, i, size)
        size[current] += size[i]


n, r, q = map(int, input().split())
graph = [[] for _ in range(n+1)]
for i in range(n-1):
    u, v = map(int, input().split())
    graph[u].append(v)
    graph[v].append(u)
tree = [[] for _ in range(n+1)]
make_tree(graph, r, -1, tree)
size = [1] * (n+1)
count_node(tree, r, size)
for i in range(q):
    u = int(input())
    print(size[u])