[BOJ][🟡5][백준#15681] 트리와 쿼리

문제 출처

BAEKJOON Online Judge #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줄에 걸쳐 각 쿼리의 답을 정수 하나로 출력한다.

예제

입력

9 5 3
1 3
4 3
5 4
5 6
6 7
2 3
9 6
6 8
5
4
8

출력

9
4
1

My Sol

import sys
input = sys.stdin.readline
sys.setrecursionlimit(200000)

def calc_subtree(u):
    global visit, rels, S
    if visit[u]: return 0
    visit[u] = 1

    cnt = 1
    for v in rels[u]:
        if visit[v]: continue
        cnt += calc_subtree(v)

    S[u] = cnt
    return cnt


V, R, Q = map(int, input().split())
rels = [[] for _ in range(V+1)]
visit = [0]*(V+1)
S = [0]*(V+1)

for _ in range(V-1):
    u, v = map(int, input().split())
    rels[u].append(v)
    rels[v].append(u)

calc_subtree(R)
for _ in range(Q):
    print(S[int(input())])

DFS와 DP를 활용해서 서브트리의 개수를 구하는 문제였다. DFS의 종료조건에 주의한다.

1. visit 배열이 1, 즉 방문한 경우 상위 노드인 것으로 간주하고 0을 반환한다.
2. leaf 노드는 상위 노드만 존재하고 이는 visit이므로 0을 반환해 1로 초기화한 cnt 자체를 반환한다. 즉, leaf 노드는 서브 트리의 노드 개수가 1개이다.
3. 상위 노드는 하위 노드 1~2개에 대해 각각의 서브트리 개수 + 자기 자신의 개수인 1을 더한 값이 서브트리 개수이다.
4. R부터 시작해 서브트리의 노드 개수를 모두 계산해주고, Q개에 대하여 출력하면 되겠다.

결과

맞았습니다!!

모범답안

출처

# empty