백준 15681: 트리와 쿼리

uni.gy·2024년 1월 31일
0

알고리즘

목록 보기
43/61

문제

풀이

  1. 트리를 구성하기 위한 간선 정보를 인접리스트에 담아둔다.
  2. 루트 노드를 시작으로 트리를 만들어준다.
  3. 서브트리 사이즈를 구해준다.
    • 처음에 매번 쿼리마다 사이즈를 구했을 때 시간초과 발생
    • 리프노드부터 부모에 자식의 서브트리 사이즈를 더해서 한번에 모든 노드의 사이즈를 구해뒀다.

코드

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

public class Main {

    static int n,r,q;
    static ArrayList<Integer>[] edges,childs;
    static int[] v,subTreeSize;

    public static void main(String[] args) throws IOException {
        // System.out.println(new Solution().solution(names, fees));
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
        StringTokenizer st;

        st=new StringTokenizer(br.readLine());
        n=Integer.parseInt(st.nextToken());
        r=Integer.parseInt(st.nextToken());
        q=Integer.parseInt(st.nextToken());
        edges=new ArrayList[n+1];
        for(int i=1;i<n+1;i++)edges[i]=new ArrayList<>();
        childs=new ArrayList[n+1];
        for(int i=1;i<n+1;i++)childs[i]=new ArrayList<>();
        v=new int[n+1];
        subTreeSize=new int[n+1];

        for(int i=0;i<n-1;i++){
            st=new StringTokenizer(br.readLine());
            int a=Integer.parseInt(st.nextToken());
            int b=Integer.parseInt(st.nextToken());
            edges[a].add(b);
            edges[b].add(a);
        }

        v[r]=1;
        makeTree(r);
        getTreeSize(r);

        for(int i=0;i<q;i++){
            int node=Integer.parseInt(br.readLine());
            bw.write(subTreeSize[node]+"\n");
        }
        bw.flush();
    }

    static void makeTree(int node){
        for(int a:edges[node]){
            if(v[a]==0){
                childs[node].add(a);
                v[a]=1;
                makeTree(a);
            }
        }
    }

    static void getTreeSize(int node){
        subTreeSize[node]=1;
        for(int a:childs[node]){
            getTreeSize(a);
            subTreeSize[node]+=subTreeSize[a];
        }
    }

}

#트리

profile
한결같이

0개의 댓글