문제풀이/백준oj

[백준OJ] 15681번 트리와 쿼리 - Java

Hyeon-Uk 2023. 1. 2. 20:26
반응형

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

 

15681번: 트리와 쿼리

트리의 정점의 수 N과 루트의 번호 R, 쿼리의 수 Q가 주어진다. (2 ≤ N ≤ 105, 1 ≤ R ≤ N, 1 ≤ Q ≤ 105) 이어 N-1줄에 걸쳐, U V의 형태로 트리에 속한 간선의 정보가 주어진다. (1 ≤ U, V ≤ N, U ≠ V)

www.acmicpc.net


풀이

먼저 그래프를 루트를 기준으로 dfs를 돌며 " 해당 노드를 루트로 하는 서브트리를 구성하는 노드의 개수" 를 구해줍니다.

문제의 예를 들어보겠습니다.

이렇게 DFS를 돌며 개수들을 모두 구해 저장한 뒤, 쿼리에 맞는 답을 출력하면 됩니다.

 

시간복잡도

DFS과정 = O(N)

쿼리 하나당 시간복잡도 = O(1)

 

따라서 모든 시간복잡도  = O(N+Q) = O(2N) = O(N)

 

코드

import java.io.*;
import java.util.ArrayList;
import java.util.StringTokenizer;

public class Main {
    static ArrayList<ArrayList<Integer>> graph = new ArrayList<>();
    static int[] nodes;
    static boolean[] visited;

    public static int dfs(int node){
        visited[node]=true;
        if(graph.get(node).isEmpty()){//리프노드
            return nodes[node]=1;
        }

        int ret = 0;
        for(int next : graph.get(node)){
            if(!visited[next])  ret+=dfs(next);
        }

        return nodes[node]=ret+1;
    }

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringBuilder sb = new StringBuilder();
        StringTokenizer st = new StringTokenizer(br.readLine());

        int N = Integer.valueOf(st.nextToken());
        int R = Integer.valueOf(st.nextToken());
        int Q = Integer.valueOf(st.nextToken());

        nodes = new int[N+1];
        visited = new boolean[N+1];
        for(int i=0;i<=N;i++) graph.add(new ArrayList<>());

        for(int i=0;i<N-1;i++){
            st = new StringTokenizer(br.readLine());
            int u = Integer.valueOf(st.nextToken());
            int v = Integer.valueOf(st.nextToken());
            graph.get(u).add(v);
            graph.get(v).add(u);
        }
        dfs(R);

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

        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
        bw.write(sb.toString());
        bw.flush();
    }
}
반응형