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

트리와 쿼리

힌트에 굉장히 좋은 코드가 있어서 따로 정리하려 한다.

루트가 없는 트리 만들기

트리의 루트 노드가 정해져 있지 않아, 하나를 루트노드로 정해서 트리를 만들고자 할때 유용한 함수이다.

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)

특정 정점을 루트로 하는 서브트리에 속한 정점의 수를 계산

계산하는 방식은 다음과 같다.

1. 먼저 V을 루트로 하는 서브트리에 포함된 노드들의 개수를 저장할 size 리스트를 만든다.
2. 현재 노드인V(루트 노드)도 해당 서브트리에 포함되므로, size에 +1을 해준다.
3. 현재 노드의 서브트리에 포함된 노드의 개수를 구하고 이를 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])