[백준/Python] BOJ 15681 - 트리와 쿼리

👀 문제

15681번: 트리와 쿼리 ✨ 골드 5

간선에 가중치와 방향성이 없는 임의의 루트 있는 트리가 주어졌을 때, 아래의 쿼리에 답해보도록 하자.

• 정점 U를 루트로 하는 서브트리에 속한 정점의 수를 출력한다.

만약 이 문제를 해결하는 데에 어려움이 있다면, 하단의 힌트에 첨부한 문서를 참고하자.

✍️ 코드

import sys
sys.setrecursionlimit(10**9)

def makeTree(currentNode, parent):
    for node in graph[currentNode]:
        if node != parent:
            # node를 currentNode의 자식으로 추가
            childs[currentNode].append(node)
            # node의 부모를 currentNode로 설정
            parents[node] = currentNode
            makeTree(node, currentNode)

def countSubtreeNodes(currentNode):
    # 자신도 자신을 루트로 하는 서브 트리에 포함되므로 0이 아닌 1에서 시작
    size[currentNode] = 1
    for node in childs[currentNode]:
        countSubtreeNodes(node)
        size[currentNode] += size[node]

# 입력 받기
N, R, Q = map(int, input().split())
graph = [[] for _ in range(N+1)]
childs = [[] for _ in range(N+1)]
parents = [0 for _ in range(N+1)]
size = [0 for _ in range(N+1)]

for _ in range(N-1):
    U, V = map(int, sys.stdin.readline().split())
    graph[U].append(V)
    graph[V].append(U)

makeTree(R, -1)
countSubtreeNodes(R)

for _ in range(Q):
    i = int(sys.stdin.readline())
    print(size[i])

📍 나는 백준의 문제에 있던 힌트를 보고 풀었는데, 그 덕분인지 생각보다 쉽게 풀렸다. 문제를 풀고 난 뒤에 구글링을 해 보니 더 간략한 풀이 방법이 있는 것 같았다. 다른 방법으로도 풀어봐야겠다.

🚑 메모리 초과 발생 원인

코드를 다 짜고 나서 백준에 돌려 보니까 메모리 초과 가 발생했다. 🤨
구글에 검색을 해 보니까 pypy3로 돌려서 메모리 초과가 발생했던 것이었다.
얌전히 python3으로 돌려주니 메모리 초과가 발생하지 않았다.