15681 - 트리와 쿼리

Byung Seon Kang·2022년 12월 3일
0

코드

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

/**
 * @Author: kbs
 */
public class Main {

    static int N, R, Q;
    static ArrayList<Integer>[] tree;
    static int[] ans;
    static boolean[] visit;

    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());

        tree = new ArrayList[N + 1];
        for (int i = 0; i <= N; i++) tree[i] = new ArrayList<>();

        ans = new int[N + 1];
        visit = new boolean[N + 1];

        Arrays.fill(ans, 1);

        for (int i = 0; i < N - 1; i++) {
            st = new StringTokenizer(br.readLine());
            int v1 = Integer.parseInt(st.nextToken());
            int v2 = Integer.parseInt(st.nextToken());
            tree[v1].add(v2);
            tree[v2].add(v1);
        }

        visit[R] = true;
        dfs(R);


        StringBuilder sb = new StringBuilder();

        for (int i = 0; i < Q; i++) {
            sb.append(ans[Integer.parseInt(br.readLine())]).append("\n");
        }

        System.out.println(sb);
    }

    private static int dfs(int root) {
        for (int next : tree[root]) {
            if (!visit[next]) {
                visit[next] = true;
                ans[root] += dfs(next);
            }
        }

        return ans[root];
    }
}
  • DFS문제
  • 서브 트리의 방문횟수를 다 더해주면 된다.
  • '정점 U를 루트로 하는 서브트리'에 주의
    • 모든 ans배열을 1로 초기화해주고 시작.
profile
왜 필요한지 질문하기

0개의 댓글