문제링크

골드 5 문제지만 스스로 다시 풀고 이해하기 위해 글을 쓴다
15681번 트리와 쿼리는 주어진 트리 구조에서 특정 정점을 루트로 하는 서브트리에 속한 노드의 개수를 구하는 문제다.

문제 개요

간선에 가중치와 방향성이 없는 임의의 루트 있는 트리가 주어졌을 때, 아래의 쿼리에 답해보도록 하자.
정점 U를 루트로 하는 서브트리에 속한 정점의 수를 출력한다.
트리의 루트는 R로 지정되며, 각 쿼리마다 해당 노드를 루트로 하는 서브트리에 속한 노드의 수를 구하는 것이 문제의 목표입니다.

접근 방식

1. 트리 구조를 먼저 구성한다
2. 서브트리 노드 개수 계산한다

코드

단방향 트리로 만드는 것이 핵심
DFS를 활용한 서브트리 크기를 계산해야한다

import java.io.*;
import java.util.*;

public class Main{
    
    static int N, R, Q; // 각각 노드의 수, 루트 노드, 쿼리의 수
    // list는 입력으로 주어진 트리의 양방향 그래프 정보 저장
    // tree는 루트 기준 구성된 단방향 트리 정보 저장
    static ArrayList<Integer>[] tree, list;
    static int[] arrQ, parent, size;
    
    public static void main(String[] args) throws Exception{
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringBuilder sb = new StringBuilder();
        StringTokenizer st = new StringTokenizer(br.readLine());
        
        N = Integer.parseInt(st.nextToken());
        R = Integer.parseInt(st.nextToken());
        Q = Integer.parseInt(st.nextToken());
        
        list = new ArrayList[N+1];
        tree = new ArrayList[N+1];
        parent = new int[N+1];
        size = new int[N+1];
        arrQ = new int[Q];
        
        for(int i=1; i<=N; i++){
            list[i] = new ArrayList<>();
            tree[i] = new ArrayList<>();
        }
        
        
        for(int i=1; i<N; i++){
            st = new StringTokenizer(br.readLine());
            int x = Integer.parseInt(st.nextToken());
            int y = Integer.parseInt(st.nextToken());
            list[x].add(y);
            list[y].add(x);
        }
        
        makeTree(R, -1);
        countSubTreeNodes(R);
        
        for(int i=0; i<Q; i++){
            int query = Integer.parseInt(br.readLine());
            sb.append(size[query]).append("\n");
        }
        
        System.out.println(sb);
    }
    
    // DFS를 사용한 루트 노드 기준으로 트리를 구성
    // 루트가 R인 단방향 트리를 만든다
    static void makeTree(int curNode, int p){
        for(int node : list[curNode]){
            if(node != p){
                tree[curNode].add(node);
                parent[node] = curNode;
                makeTree(node, curNode);
            }
        }
    }
    
    // 서브트리 크기 계산
    static void countSubTreeNodes(int curNode){
        size[curNode] = 1;
        for(int node : tree[curNode]){
            countSubTreeNodes(node);
            size[curNode] += size[node];
        }
    }
}