[BOJ] 15681번: 트리와 쿼리 (Subtree, DFS)

 

(( 시간 효율 지키면서 서브트리 구별하는 방법 외우기 ))

 

문제 링크

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

 

 

 

TIL

• 2 ≤ N ≤ 100,000 , 1 ≤ R ≤ N, 1 ≤ Q ≤ 100,000
  • 이렇게 입력이 10만을 초과하는 경우 무조건 sys.stdin으로 input 받기 !!!!!!!
  • 그냥 input()으로 받으면 100% 시간 초과 발생

 

 

[ 처음 코드 ] - 시간 초과

from collections import deque, defaultdict

N, R, Q = map(int, input().split())

# Level Tree 만들기
dd = defaultdict(list)
level = [-1]*(N+1)
level[R] = 0

for i in range(N-1):
    u, v = map(int, input().split())
    dd[u].append(v)
    dd[v].append(u)

q = deque()
q.append((R,0))
while q:
    node, lev = q.popleft()
    for n in dd[node]:
        if level[n] < 0 :
            level[n] = lev+1
            q.append((n,lev+1))

# 순회
for i in range(Q):
    s = int(input())

    treeq = deque()
    visited = set()
    treeq.append(s)
    visited.add(s)
    cnt = 0

    while treeq:
        tree_node = treeq.popleft()
        cnt += 1

        for neighbor in dd[tree_node]:
            if level[neighbor] > level[tree_node] and neighbor not in visited:
                treeq.append(neighbor)
    print(cnt)

• 시간 효율성이 너무 떨어짐
  • 쿼리마다 BFS → 최악의 경우 O(N)
  • 쿼리 수가 Q=100,000일 때, 시간복잡도: O(N×Q) = 최대 10¹⁰
  • → 시간 초과(TLE) 가능성 매우 높음

 

 

 

Code

from collections import defaultdict
import sys
sys.setrecursionlimit(10**6)    # 10**5는 Recursionlimit error 발생
input = sys.stdin.readline  # 이거 없으면 시과초과 발생 !! 

N, R, Q = map(int, input().split())
graph = defaultdict(list)
for i in range(N-1):
    u, v = map(int,input().split())
    graph[v].append(u)
    graph[u].append(v)

subtree_size = [0]*(N+1)
visited = [False]*(N+1)

def dfs(node):
    visited[node] = True
    size = 1
    
    for neighbor in graph[node]:
        if not visited[neighbor]:
            size += dfs(neighbor)
    subtree_size[node] = size
    return size

dfs(R)
for i in range(Q):
    print(subtree_size[int(input())])

• 최대 길이(깊이) 구하는 dfs랑 비슷한 원리지만, size = 1 로 항상 초기화한다는 것만 다름 !! 
  • 본인을 포함해서 subtree의 개수를 세기 때문에 1로 초기화
• 그냥 한 번에 전체 서브트리 다 구해놓고 사용하는 방법

 

두 코드의 시간복잡도 비교

방식	탐색	쿼리	전체 시간복잡도
1번 방식 (Level 모두 순회)	O(N) x Q	O(N)	O(NQ)      ❌
2번 방식 (DFS 선처리 방식)	O(N)	O(1) x Q	O(N + Q)  ✅