[dfs/백준] 15681번 : 트리와 쿼리

- 문제 : https://www.acmicpc.net/problem/15681

 

Gold Level5 문제 

문제에 있는 설명을 기준으로 풀면 된다고해서 최대한 따라해보려고 했는데, 55%에서 시간초과가 발생해서 이전 풀이를 보고 진행했다 

 

주어진 노드마다 그걸 루트노드로 한 뒤 자식 노드들의 갯수를 세는 방법은 시간초과가 발생한다

그래서 주어진 루트노드를 중심으로 각 노드가 루트노드인 서브 트리의 노드 갯수를 dp에 저장한 뒤 필요할 때마다 출력하는 방식으로 진행하면 된다 

 

본 트리 하나 만든 후, 각 노드의 자식노드 갯수를 세는 함수를 따로 만들어야하는 거라고 생각했는데, 그게 아니라 루트를 만들어가면서 dp리스트에 값을 추가하는 것이었다 

흠,, 역시 dfs는 쉽지 않은 것 같다 어렵!

 

정답풀이에서 dfs는 잎 노드일 때까지 내려간 다음, 그것의 부모노드에 자식 노드의 갯수를 하나씩 추가해간다. 

그 부모노드의 자식노드는 자식노드의 서브트리 노드 갯수를 누적해가는 방식이다 

 

- 정답풀이 :

import sys
sys.setrecursionlimit(10**6)
input = sys.stdin.readline #이걸 설정해줘야 시간초과 발생하지 않는다 

n,r,q = map(int,input().split())
data = [[] for _ in range(n+1)]
dp = [0]*(n+1)

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

def dfs(node):
    dp[node] = 1
    for i in data[node]:
        if not dp[i]:
            dfs(i)
            #잎노드에서부터 하나씩 추가해 가야하므로
            #dfs 후에 dp값을 누적해간다 
            dp[node] += dp[i]  
        
dfs(r)
for _ in range(q):   
    print(dp[int(input())])

 

- 시도해본 풀이 

55%에서 시간초과가 발생했다 

makeTree()함수 때문에 초기값 입력 받는 그래프 따로, 자식노드만 입력하는 트리를 따로 만들어야 한다고 생각해서

data, graph 리스트를 각각 만들었다 

 

import sys
input = sys.stdin.readline

n,r,q = map(int,input().split())
data = [[] for _ in range(n+1)]
graph = [[] for _ in range(n+1)]

for _ in range(n-1):
    a,b = map(int,input().split())
    #부모노드, 자식노드 구별이 안됨 
    data[a].append(b)
    data[b].append(a)

def makeTree(currentNode, parent):
    for i in data[currentNode]:
        if i != parent:
            graph[currentNode].append(i)
            makeTree(i,currentNode)
            
makeTree(r,-1)

def countSubTreeNodes(node):
    visit[node] = 1
    #root의 자식 노드들 탐색
    for i in graph[node]:
        countSubTreeNodes(i)
        visit[node] += visit[i]
            
visit = [0]*(n+1)    
for _ in range(q):
    x = int(input())    
    countSubTreeNodes(r)
    print(visit[x])