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

 

 

문제해결방법 - 

 

1. dfs(깊이 우선 탐색), bfs(너비 우선 탐색) 알고리즘을 이용할 수 있다.

1). bfs의 경우, dfs의 재귀 함수처럼 맨 끝 노드부터 누적되는 구조를 활용할 수 없기 때문에, 다른 장치가 필요하다.

2. dfs로 끝 노드까지 재귀 함수가 실행되고, 방문리스트에 끝 노드의 갯수가 방문리스트[부모노드]에 더해질 수 있도록 구성한다.

3. 그 다음 q줄에 주어지는 U의 값(U를 루트노드로 하는 트리의 갯수)을 출력한다.

 

- dfs(깊이 우선 탐색)

import sys
from collections import deque
import heapq
sys.setrecursionlimit(10**6)
input = sys.stdin.readline

n, r, q = map(int, input().split())
graph = [[] for _ in range(n+1)]
v = [0 for _ in range(n+1)]
for _ in range(n-1):
	a, b = map(int, input().split())
	graph[a].append(b)
	graph[b].append(a)

def dfs(idx):
	v[idx] = 1
	for next in graph[idx]:
		if v[next] == 0:
			dfs(next)
			v[idx] += v[next]
dfs(r)
for i in range(q):
	root = int(input())
	print(v[root])

 

 

- bfs(너비 우선 탐색)

import sys
from collections import deque
import heapq
sys.setrecursionlimit(10**6)
input = sys.stdin.readline

n, r, q = map(int, input().split())
graph = [[] for _ in range(n+1)]
v = [0 for _ in range(n+1)]
root = [0 for _ in range(n+1)]
dp = [1 for _ in range(n+1)]
for _ in range(n-1):
	a, b = map(int, input().split())
	graph[a].append(b)
	graph[b].append(a)

def bfs(idx):
	q = deque([idx])
	v[idx] = 1
	temp = []
	while q:
		idx = q.popleft()
		temp.append(idx)
		for next in graph[idx]:
			if v[next] == 0:
				v[next] = 1
				root[next] = idx
				q.append(next)
	for i in temp[::-1]:
		for j in graph[i]:
			if j != root[i]:
				dp[i] += dp[j]
	return
			
bfs(r)
for i in range(q):
	root = int(input())
	print(dp[root])