[15681] 트리와 쿼리

🔗 문제 링크

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

🔍 문제 설명

간선에 가중치와 방향성이 없는 임의의 루트 있는 트리가 주어졌을 때, 아래의 쿼리에 답해보도록 하자.

정점 U를 루트로 하는 서브트리에 속한 정점의 수를 출력한다.
만약 이 문제를 해결하는 데에 어려움이 있다면, 하단의 힌트에 첨부한 문서를 참고하자.

⚠️ 제한사항

• $(2 ≤ N ≤ 105, 1 ≤ R ≤ N, 1 ≤ Q ≤ 105)$

• $(1 ≤ U, V ≤ N, U ≠ V)$

• $(1 ≤ U ≤ N)$

• 입력으로 주어지는 트리는 항상 올바른 트리임이 보장된다.

🗝 풀이 (언어 : Java)

DFS와 DP를 동시에 이용해서 푸는 문제이다. 리스트와 배열을 이용해 양방향 그래프를 만들어주고, DFS로 루트부터 말단 노드까지 재귀로 자식을 탐색한다. 이과정에서 현재 노드의 숫자 인덱스마다 1씩 값을 준다.(개수) 그리고 재귀가 끝나면 자식 숫자의 인덱스에 저장된 값을 부모 노드 인덱스의 값에 합해준다. 이렇게 부모까지 쭉 해주면 모든 경우들의 정답이 저장된 배열을 만들 수 있다.

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

class Main {
    private static int[] dp;
    private static void dfs(int x, int prev, List<Integer>[] arr) {
        dp[x] = 1; // 현재 위치 정점 카운트
        for(Integer next : arr[x]) {
            if(next.intValue() == prev) // 연결된 정점 중에 부모는 스킵
                continue;
            dfs(next, x, arr); // 자식 정점으로 탐색
            dp[x] += dp[next]; // 하위 그래프의 정점들 수를 부모에 더해주기
        }
    }

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine(), " ");
        int n = Integer.parseInt(st.nextToken()), r = Integer.parseInt(st.nextToken()), q = Integer.parseInt(st.nextToken());
        List<Integer>[] arr = new ArrayList[n+1];
        for(int i = 1; i <= n; i++)
            arr[i] = new ArrayList<>();
        dp = new int[n+1];
        for (int i = 1; i < n; i++) {
            StringTokenizer st2 = new StringTokenizer(br.readLine());
            int u = Integer.parseInt(st2.nextToken()), v = Integer.parseInt(st2.nextToken());
            arr[u].add(v); // 양방향이므로 둘다 반대쪽 노드 넣기
            arr[v].add(u);
        }
        dfs(r, -1, arr); // DFS 탐색으로 DP 배열 완성
        for (int i = 0; i < q; i++)
            System.out.println(dp[Integer.parseInt(br.readLine())]);
        br.close();
    }
}