15681. 트리와 쿼리

트리에서의 동적계획법.

재귀 호출을 통해서 노드의 개수를 세줄 수 있다.

import sys;input=sys.stdin.readline
sys.setrecursionlimit(10**9)
n,r,q=map(int,input().split())
tree = [[] for _ in range(n+1)]
cnt = [0] * (n+1)

def NodeCount(x):
  cnt[x] = 1
  for i in tree[x]:
    if not cnt[i]:
      NodeCount(i)
      cnt[x] += cnt[i]

for _ in range(n-1):
  u, v = map(int,input().split())
  tree[u].append(v)
  tree[v].append(u)
NodeCount(r)

for i in range(q):
  temp = int(input())
  print(cnt[temp])

백준은 자바스크립트의 콜스택 사이즈가 좀 작은것? 같다. 파이썬코드를 자스로 옮겼을 때 틀렸다고 나온다.

const input = require("fs")
  .readFileSync("./dev/stdin")
  .toString()
  .trim()
  .split("\n")

const [n, r, q] = input[0].split(" ").map(Number)
const edges = input
  .slice(1, input.length - q)
  .map((x) => x.split(" ").map(Number))
const queries = input.slice(input.length - q).map(Number)
const tree = Array.from({ length: n + 1 }, () => [])
const cnt = Array.from({ length: n + 1 }, () => 0)
for (let [a, b] of edges) {
  tree[a].push(b)
  tree[b].push(a)
}
function dfs(x) {
  cnt[x] = 1
  for (let y of tree[x]) {
    dfs(y)
    cnt[x] += cnt[y]
  }
}
dfs(r)
queries.map((q) => console.log(cnt[q]))