[자바] 백준 15681 - 트리와 쿼리 (java)

목차

문제 : boj15681

 

 

필요 알고리즘

• DFS, 트리DP
  
  • DFS로 트리DP를 진행하는 문제이다.

※ 제 코드에서 왜 main 함수에 로직을 직접 작성하지 않았는지, 왜 Scanner를 쓰지 않고 BufferedReader를 사용했는지 등에 대해서는 '자바로 백준 풀 때의 팁 및 주의점' 글을 참고해주세요. 백준을 자바로 풀어보려고 시작하시는 분이나, 백준에서 자바로 풀 때의 팁을 원하시는 분들도 보시는걸 추천드립니다.

 

 

풀이

  우선 간선들을 입력받고, 루트가 주어져 있으므로 루트부터 DFS를 진행한다.

이 때 dfs() 함수는 리턴으로 자기자신을 포함해 자신 이하의 서브트리의 갯수를 리턴하는 것으로 정의한다.

그럼 이하처럼 코드를 짤 수 있고, 루트부터 시작해 모든 정점에 대해 DFS 진행하면서 dfs() 함수의 리턴값을 메모이제이션 해두면, 이후 Q개의 쿼리에 대해 각 O(1)로 처리 가능하다. dfs() 코드는 이하와 같다.

private int dfs(final int r) {
    int cnt = 1;

    for (Integer next : edges[r]) {
        if (v[next]) continue;
        v[next] = true;

        cnt += dfs(next);
    }

    return answer[r] = cnt;
}

 

 

코드 : github

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

public class Main {
    static BufferedReader br = new BufferedReader(new InputStreamReader(System.in));

    public static void main(String[] args) throws Exception {
        new Main().solution();
    }

    List<Integer>[] edges;
    int[] answer;
    boolean[] v;

    private void solution() throws Exception {
        StringTokenizer st = new StringTokenizer(br.readLine());
        int n = Integer.parseInt(st.nextToken());
        int r = Integer.parseInt(st.nextToken());
        int q = Integer.parseInt(st.nextToken());

        answer = new int[n+1];
        v = new boolean[n+1];
        edges = new List[n+1];
        for (int i = 1; i <= n; i++) edges[i] = new ArrayList<>();
        for (int i = 0; i < n-1; i++) {
            st = new StringTokenizer(br.readLine());
            int u = Integer.parseInt(st.nextToken());
            int v = Integer.parseInt(st.nextToken());
            edges[u].add(v);
            edges[v].add(u);
        }

        v[r] = true;
        dfs(r);

        StringBuilder sb = new StringBuilder();
        while (q-->0) {
            int u = Integer.parseInt(br.readLine());

            sb.append(answer[u]).append('\n');
        }

        System.out.print(sb);
    }

    private int dfs(final int r) {
        int cnt = 1;

        for (Integer next : edges[r]) {
            if (v[next]) continue;
            v[next] = true;

            cnt += dfs(next);
        }

        return answer[r] = cnt;
    }
}