[Python] 백준 15681 트리와 쿼리

문제

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

• 정점 U를 루트로 하는 서브트리에 속한 정점의 수를 출력한다.

만약 이 문제를 해결하는 데에 어려움이 있다면, 하단의 힌트에 첨부한 문서를 참고하자.

https://www.acmicpc.net/problem/15681

 

 

 

난이도

골드 5

 

 

 

내 코드

import sys
sys.setrecursionlimit(10**6)

def DFS(currentNode) :
    visit[currentNode] = True 

    for nextNode in graph[currentNode] :
        if not visit[nextNode] :
            answer[currentNode] += DFS(nextNode)
    
    return answer[currentNode]


n, r, q = map(int,sys.stdin.readline().split())
graph = [[] for _ in range(n+1)]
visit = [False] * (n+1)
answer = [1] * (n+1)
for _ in range(n-1) : 
    u, v = map(int,sys.stdin.readline().split())
    graph[u].append(v)
    graph[v].append(u)

DFS(r)

for _ in range(q) :
    x = int(sys.stdin.readline())
    print(answer[x])

 

 

 

 

 

 

 

풀이과정

DFS로 풀려고 했다. 처음에는 스택을 이용해서 DFS 순환을 하면서 pop() 이 일어나는 순간에 현재 노드의 값을 부모노드에게 더해주는 방법을 사용했는데, 답은 맞는거 같은데 65% 에서 시간초과가 발생했다. 

스택을 이용하는 과정에서 pop() 과 append()가 발생하고 answer[] 배열도 계속 업데이트 하는 과정에서 시간을 오래걸린거 같았다.

그래서 재귀함수를 이용한 DFS로 방법을 바꾸었다.

현재 노드에서 연결된 자식 노드로 내려가면 서브트리 크기를 구하고 부모 노드에 더해주는 방식은 비슷하지만 재귀 방식에서의 메모리 접슨 효율성 차이로 시간초과를 해결할 수 있었다.