힌트에 굉장히 좋은 코드가 있어서 따로 정리하려 한다.
트리의 루트 노드가 정해져 있지 않아, 하나를 루트노드로 정해서 트리를 만들고자 할때 유용한 함수이다.
makeTree(5, -1)
을 실행하게 되면 root가 5인 트리를 만들 수 있게 된다. -1은 부모가 없음을 의미한다.
def makeTree(currentNode, parent) :
for(Node in connect[currentNode]) :
if Node != parent:
add Node to currentNode’s child
set Node’s parent to currentNode
makeTree(Node, currentNode)
계산하는 방식은 다음과 같다.
- 먼저
V
을 루트로 하는 서브트리에 포함된 노드들의 개수를 저장할size
리스트를 만든다.- 현재 노드인
V
(루트 노드)도 해당 서브트리에 포함되므로,size
에 +1을 해준다.- 현재 노드의 서브트리에 포함된 노드의 개수를 구하고 이를
size
에 더해준다.
만약 자식이 없다면, 즉 리프노드라면 child가 없고 서브트리에 본인만 있기 때문에 size는 1이다. 재귀적으로 다시 돌아가며 서브트리의 노드의 개수를 구할 수 있다.
def countSubtreeNodes(currentNode) :
size[currentNode] = 1 // 자신도 자신을 루트로 하는 서브트리에 포함되므로 0이 아닌 1에서 시작한다.
for Node in currentNode’s child:
countSubtreeNode(Node)
size[currentNode] += size[Node]
처음 주어진 root를 기준으로 트리를 만들고, 각 노드의 서브트리의 개수를 미리 세는 방식을 선택했다.
import sys
sys.setrecursionlimit(10**6)
input = sys.stdin.readline
N, R, Q = map(int, input().split()) # 트리의 정점의 수, 루트의 번호, 쿼리의 수
adjM = [[] for _ in range(N + 1)] # 간선 정보를 저장할 행렬
children = [[] for _ in range(N + 1)] # 자식 정보를 저장할 리스트
size = [0] * (N + 1)
for _ in range(N - 1):
e1, e2 = map(int, input().split())
adjM[e1].append(e2)
adjM[e2].append(e1)
def make_tree(current, parent):
for near in adjM[current]:
if near != parent:
children[current].append(near)
make_tree(near, current)
def count_sub(current):
size[current] = 1
for child in children[current]:
count_sub(child)
size[current] += size[child]
make_tree(R, -1)
count_sub(R)
for _ in range(Q): # 각 쿼리에 대한 답 찾기
root = int(input())
print(size[root])