[BOJ] 15681. 트리와 쿼리

[작성일]

• 2022-03-02

[분류]

• 트리
• dp
• dfs
  
  
  

[문제링크]

• https://www.acmicpc.net/problem/15681
  
  
  

[요구사항]

• 주어진 노드를 루트로 하는 서브트리의 노드 갯수를 구하라.
  
  
  

[풀이]

문제에 구현방법 설명이 아주 상세하게 나와있어서 따로 알고리즘이 없다고 생각하고 풀었는데 인터넷을 찾아보니 더 심플하게 dp 개념을 이용해서 코드를 작성할 수 있었다.

이 문제는 tree + dp 문제이다.
우선 그래프를 입력받는다.
이후 tree를 재귀적으로 만들어가면서 해당 노드를 root로 하는 서브트리의 노드 갯수를 dp배열에 저장해서 size를 모두 구한다.

[코드 - 문제설명 그대로 따라 푼 것]

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.StringTokenizer;

public class Main {
    static int N, R, Q;
    static Node[] node;
    static ArrayList<Integer>[] connect;
    static int[] size;

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());
        N = Integer.parseInt(st.nextToken());
        R = Integer.parseInt(st.nextToken());
        Q = Integer.parseInt(st.nextToken());

        node = new Node[N + 1];
        for (int i = 1; i <= N; i++) {
            node[i] = new Node();
        }
        connect = new ArrayList[N + 1];
        for (int i = 1; i <= N; i++) {
            connect[i] = new ArrayList<>();
        }
        int u, v;
        for (int i = 0; i < N - 1; i++) {
            st = new StringTokenizer(br.readLine());
            u = Integer.parseInt(st.nextToken());
            v = Integer.parseInt(st.nextToken());

            connect[u].add(v);
            connect[v].add(u);
        }

        // makeTree
        makeTree(R, -1);

        size = new int[N + 1];
        countSubtreeNodes(R);

        StringBuilder resSb = new StringBuilder();
        for (int i = 0; i < Q; i++) {
            int idx = Integer.parseInt(br.readLine());
            resSb.append(size[idx] + "\n");
        }
        System.out.println(resSb);
    }

    public static void makeTree(int currentNode, int parent) {
        for (int v : connect[currentNode]) {
            if (v != parent) {
                node[currentNode].child.add(v);
                node[v].parent = currentNode;
                makeTree(v, currentNode);
            }
        }
    }

    public static void countSubtreeNodes(int currentNode) {
        size[currentNode] = 1;
        for (int v : node[currentNode].child) {
            countSubtreeNodes(v);
            size[currentNode] += size[v];
        }
    }
}

class Node {
    ArrayList<Integer> child = new ArrayList<>();
    int parent;
}

[코드 - tree + dp]

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.StringTokenizer;

public class Main {
    static int N, R, Q;
    static ArrayList<Integer>[] list;
    static int[] dp;

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        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];
        for (int i = 1; i <= N; i++) {
            list[i] = new ArrayList<>();
        }
        int u, v;
        for (int i = 0; i < N - 1; i++) {
            st = new StringTokenizer(br.readLine());
            u = Integer.parseInt(st.nextToken());
            v = Integer.parseInt(st.nextToken());

            list[u].add(v);
            list[v].add(u);
        }

        // makeTree
        dp = new int[N + 1];
        makeTree(R, -1);

        // print
        StringBuilder res = new StringBuilder();
        for(int i = 0; i < Q; i++) {
            res.append(dp[Integer.parseInt(br.readLine())] + "\n");
        }
        System.out.println(res);
    }

    public static int makeTree(int curr, int parent) {
        dp[curr] = 1;

        for (int child : list[curr]) {
            if(parent != child) {
                dp[curr] += makeTree(child, curr);
            }
        }

        return dp[curr];
    }
}