[BOJ] 15681 트리와 쿼리 바로가기
간선에 가중치와 방향성이 없는 임의의 루트 있는 트리가 주어졌을 때, 아래의 쿼리에 답해보도록 하자.
만약 이 문제를 해결하는 데에 어려움이 있다면, 하단의 힌트에 첨부한 문서를 참고하자.
트리의 정점의 수 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줄에 걸쳐 각 쿼리의 답을 정수 하나로 출력한다.
Union & Find
알고리즘을 이용한 방법으로 문제에 접근하게 되었다.하단의 힌트
에서 주어진 makeTree()
함수와 countSubtreeNodes()
함수를 코드에 적용하여 문제를 해결하는 방식이었다.defaultdict(list)
로 정의한 connect
딕셔너리를 생성하여 각 정점의 연결 관계 정보를 저장한다.connect = defaultdict(list)
for u, v in UV:
connect[u].append(v)
connect[v].append(u)
currentNode
를 루트로 하는 서브트리를 생성한다.currentNode
에 연결된 모든 자식 정점들을 순회한다.Node
)과 부모 정점(parent
)이 같지 않다면,tree
의 currentNode
리스트에 현재 정점(Node
)를 추가한다.Node
)의 부모를 currentNode
로 설정한다.Node
)을 루트로 하는 makeTree()
함수를 실행한다.def makeTree(currentNode, parent) :
for Node in connect[currentNode] :
if Node != parent:
#add Node to currentNode’s child
tree[currentNode].append(Node)
#set Node’s parent to currentNode
parents[Node] = currentNode
makeTree(Node, currentNode)
currentNode
)을 루트로 하는 서브트리에 속한 정점의 수를 센다.Node
)를 인자로 countSubtreeNodes()
함수를 실행한다.def countSubtreeNodes(currentNode) :
size[currentNode] = 1 # 자신도 자신을 루트로 하는 서브트리에 포함되므로 0이 아닌 1에서 시작한다.
for Node in tree[currentNode]:
countSubtreeNodes(Node)
size[currentNode] += size[Node]
from sys import stdin, setrecursionlimit
from collections import defaultdict
setrecursionlimit(10**6) # 최대 재귀 깊이의 값을 10^6 으로 재설정
def solution(N, R, UV, U):
connect = defaultdict(list)
parents = list(range(N+1))
tree = defaultdict(list)
size = list(0 for _ in range(N+1))
# 연결 관계 초기화
for u, v in UV:
connect[u].append(v)
connect[v].append(u)
# 각 정점을 루트로 하는 서브트리에 속한 정점의 수 count
def countSubtreeNodes(currentNode) :
size[currentNode] = 1 # 자신도 자신을 루트로 하는 서브트리에 포함되므로 0이 아닌 1에서 시작한다.
for Node in tree[currentNode]:
countSubtreeNodes(Node)
size[currentNode] += size[Node]
# R을 정점으로 하는 트리 생성
def makeTree(currentNode, parent) :
for Node in connect[currentNode] :
if Node != parent:
#add Node to currentNode’s child
tree[currentNode].append(Node)
#set Node’s parent to currentNode
parents[Node] = currentNode
makeTree(Node, currentNode)
makeTree(R,-1)
countSubtreeNodes(R)
# 정점 u를 루트로 하는 서브트리에 속한 정점의 수 출력
for u in U:
print(size[u])
N, R, Q = map(int,stdin.readline().split())
UV = list(list(map(int,stdin.readline().split())) for _ in range(N-1))
U = list(int(stdin.readline()) for _ in range(Q))
solution(N, R, UV, U)