[boj] 15681. 트리와 쿼리 (node.js)

문제 요약

[boj] 15681. 트리와 쿼리 (node.js)

풀이

• 주어진 무방향 트리에서 '입력된 정점을 루트로 하는 서브트리에 속한 정점의 수를 출력한다' 는 쿼리가 주어질 때, 이 쿼리를 만족하는 결과를 구현하는 문제
• 루트 노드가 주어진다.

내 풀이

• 주어진 트리를 입력받은 후, 루트 노드에서 dfs 를 구현한다.
• 이때 dfs 함수는 노드에 자식 노드가 있는 경우, 자식 노드에서 다시 dfs를 수행한다.
  • dfs 함수를 호출할 때마다 해당 노드 이하 노드의 개수를 반환하고, 반환한 값을 계속 누적 합산한다.
  • 하위 노드에 대해 누적된 값을 모두 합산하고 나면, 결과값을 result 배열의 노드 인덱스에 저장해 두어 쿼리 호출시 배열 값을 반환한다.

const readline = require("readline");
const rl = readline.createInterface({
  input: process.stdin,
  output: process.stdout,
});

const solution = () => {
  const [N, R, Q] = input().split(" ").map(Number);
  let tree = [];
  for (let i = 0; i < N - 1; i++) {
    let arr = input().split(" ").map(Number);
    if (!tree[arr[0]]) tree[arr[0]] = [];
    if (!tree[arr[1]]) tree[arr[1]] = [];
    tree[arr[0]].push(arr[1]);
    tree[arr[1]].push(arr[0]);
  }

  let result = [];
  let visited = [];
  visited[R] = 1;
  result[R] = 0;
  dfs(R, 0);

  let results = [];
  for (let i = 0; i < Q; i++) {
    let node = input();
    results.push(result[node]);
  }
  console.log(results.join("\n"));

  function dfs(node) {
    let cnt = 1;
    tree[node].forEach((next) => {
      if (visited[next]) return;
      visited[next] = 1;
      let sum = dfs(next);
      cnt += sum;
    });
    if (!result[node]) result[node] = 0;
    result[node] += cnt;
    return cnt;
  }
};

let _line = 0;
const input = () => stdin[_line++];

let stdin = [];
rl.on("line", function (line) {
  stdin.push(line);
}).on("close", function () {
  solution();
  process.exit();
});

주절주절

• 초반에 누적값을 합산하지 않고 배열에 따로따로 저장하려고 하느라 시간이 오래 걸렸다. 함수 호출한 후 반환되는 값을 누적하여 그 결과값을 사용할 수 있음을 잊지 말자!