15681_트리와 쿼리

▶ 골드5 (DFS)

문제

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

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

---

풀이

DFS를 활용하여 자식 노드의 개수를 세는 문제였습니다.

 

처음에 자식 노드 개수 세는 함수 세우고 쿼리 for문 돌려서 매번 호출했더니 시간 초과가 발생하였습니다.

 

따라서 매번 세는 방식 말고, 미리 자식 노드의 개수를 배열에 담아 놓는 방식을 생각했습니다.

DFS로 현재 노드의 자식 노드들의 개수의 합을 재귀를 통해 구현했습니다.

하지만 잘못 생각한게, 트리가 무조건 작은 수부터 큰 수로 순차적으로 구성된다고 생각해서

루트와의 거리를 비교하여 가까운 노드를 부모로, 먼 노드를 자식으로 하여 단방향 간선만 넣었습니다.

 

때문에 트리가 잘못 구성되는 경우가 발생하여 틀렸고, 이후 양방향으로 수정하고 visited 배열로 방문 여부를 추가하였습니다.

 

최종 풀이는,

1) 간선들을 입력받고 양방향으로 dictionary 형태로 넣어줍니다.

2) 방문 여부를 체크하는 자식 노드 개수 카운트 함수를 실행하여, 모든 노드들의 자식 노드 수를 배열에 담아둡니다.

3) 쿼리를 for문 돌려 각각의 값을 print 합니다.

코드

# 15681_트리와 쿼리
import sys
sys.setrecursionlimit(10**6)
input = sys.stdin.readline

n, r, q = map(int, input().split())
edge = [list(map(int, input().split())) for _ in range(n - 1)]
queries = [int(input()) for _ in range(q)]

visited = [False for _ in range(n)]
count = [0 for _ in range(n)] # 자식 노드 개수
graph = {} # 그래프

for i in range(n):
    graph[i + 1] = []

for u, v in edge:
    graph[u].append(v)
    graph[v].append(u)

# 재귀를 통해서 모든 노드의 자식 노드의 개수를 구함
def countNode(cur):
    visited[cur - 1] = True
    if len(graph[cur]) == 0:
        count[cur - 1] = 1
        return 1

    count[cur - 1] += 1
    for child in graph[cur]:
        if not visited[child - 1]:
            count[cur - 1] += countNode(child)

    return count[cur - 1]

countNode(r)
for que in queries:
    print(count[que - 1])