[백준] 15681번 : 트리와 쿼리 - 자바(Java)

Problem 🔒

문제

(15681번 : 트리와 쿼리)

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

• 정점 U를 루트로 하는 서브트리에 속한 정점의 수를 출력한다.

만약 이 문제를 해결하는 데에 어려움이 있다면, 하단의 힌트에 첨부한 문서를 참고하자.

입력

트리의 정점의 수 N과 루트의 번호 R, 쿼리의 수 Q가 주어진다. (2 ≤ N ≤ 105, 1 ≤ R ≤ N, 1 ≤ Q ≤ 105)

이어 N-1줄에 걸쳐, U V의 형태로 트리에 속한 간선의 정보가 주어진다. (1 ≤ U, V ≤ N, U ≠ V)

이는 U와 V를 양 끝점으로 하는 간선이 트리에 속함을 의미한다.

이어 Q줄에 걸쳐, 문제에 설명한 U가 하나씩 주어진다. (1 ≤ U ≤ N)

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

출력

Q줄에 걸쳐 각 쿼리의 답을 정수 하나로 출력한다.

예제 입력 1

9 5 3
1 3
4 3
5 4
5 6
6 7
2 3
9 6
6 8
5
4
8

예제 출력 1

9
4
1

---

Approach ⭕

이 문제는 특정 노드를 루트로 하는 서브트리에 속한 정점의 수를 구해야한다.

서브트리는 특정 노드를 포함한 하위 노드들로 구성된다.

DFS를 사용하여 각 노드를 탐색하면서 서브트리의 정점 개수를 계산하여 배열에 저장했다.

1. 트리 생성
   • 주어지는 트리는 사이클이 없는 무방향 그래프이다.
   • 그렇기에, 간선의 정보를 양방향으로 저장한다.
2. DFS 초기화
   • 각 노드를 루트로 하는 서브트리의 정점 개수를 dp[]에 저장한다.
   • 모든 노드의 서브트리는 자신을 포함하므로 dp[]를 1로 초기화한다.
   • 정점의 방문 여부를 visit[]에 저장한다.
3. DFS 탐색(순서 중요⭐)
   1. 현재 노드를 방문 처리한다.
   2. 현재 노드와 연결된 자식 노드를 탐색한다.
      • 연결된 정점들 중 방문한 정점은 부모 정점밖에 없다.
      • 방문한 정점은 방문하지 않는다.
   3. 모든 자식 노드 탐색이 끝난 후, 부모 노드에 서브트리 크기를 더해준다.

---

Solution 💡

public class 트리와쿼리 {
    static int N, R, Q, u, v, dp[], U;
    static boolean[] visit;
    static ArrayList<ArrayList<Integer>> list = new ArrayList<>();
    static StringBuilder sb = new StringBuilder();
    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()); // 쿼리
        visit = new boolean[N+1]; // 방문 배열
        dp = new int[N+1]; // 서브트리에 속한 정점 수 저장할 배열
        Arrays.fill(dp,1); // 자신을 포함하는 의미로 1로 초기화
        for(int i=0; i<=N; i++){list.add(new ArrayList<>());}
        for(int i=0; i<N-1; i++){
            st = new StringTokenizer(br.readLine());
            u = Integer.parseInt(st.nextToken());
            v = Integer.parseInt(st.nextToken());
            // 무방향 그래프이므로 양방향으로 간선 저장
            list.get(u).add(v);
            list.get(v).add(u);
        }
        DFS(R,0);
        for(int i=0; i<Q; i++){
            U = Integer.parseInt(br.readLine());
            sb.append(dp[U]+"\\n");
        }
        System.out.println(sb);
    }
    public static void DFS(int cur, int parent){
        // 1. 현재 노드 방문처리
        visit[cur] = true;
        // 2. 현재 노드와 연결된 자식 노드 탐색
        for(int nxt : list.get(cur)){
            if(visit[nxt]) continue;
            DFS(nxt, cur);
        }
        // 3. 자식 노드 탐색 후, 부모 노드 갱신
        dp[parent] += dp[cur];
    }
}