[Algorithm] 백준 트리 DP문제 - 트리와 쿼리(15861) 자바 풀이

안녕하세요. 회사와 함께 성장하고 싶은 KOSE입니다.

이번 포스팅은 백준 트리 DP 문제 트리와 쿼리 자바 풀이를 진행하도록 하겠습니다.

문제 출처: https://www.acmicpc.net/problem/15681

15681번: 트리와 쿼리

 

1. 풀이 소스

 

import java.io.*;
import java.util.*;
import static java.lang.Integer.parseInt;

public class Main {
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));    
        StringTokenizer st = new StringTokenizer(br.readLine());
        
        int n = parseInt(st.nextToken());
        int r = parseInt(st.nextToken());
        int q = parseInt(st.nextToken());
        
        Graph graph = new Graph(n, r);
        for (int i = 0; i < n - 1; i++) {
            st = new StringTokenizer(br.readLine());
            graph.addEdge(parseInt(st.nextToken()), parseInt(st.nextToken()));           
        }
        
        for (int i = 0; i < q; i++) graph.addQuery(parseInt(br.readLine()));
        
        graph.makeSubTree(r); // 서브 트리 만들며 dp 저장
        graph.getPrint(); // 결과 출력
    }
}

class Graph {
    int n; // n개의 노드
    int root; // 루트 노드
    Node[] nodes; // 노드 저장
    List<Integer> querys = new ArrayList<>(); // 출력해야하는 쿼리
    int[] dp; // 서브 트리 개수
    
    Graph (int n, int root) {
        this.n = n;
        this.root = root;
        nodes = new Node[n + 1];
        dp = new int[n + 1];
        
        for (int i = 1; i <= n; i++) nodes[i] = new Node(i);
    }
    
    class Node {
        int x;
        List<Node> adjacent = new ArrayList<>(); // 인접 노드
        
        Node (int x) {
            this.x = x;
        }
    }
    
    void addEdge(int x, int y) {
        Node n1 = nodes[x];
        Node n2 = nodes[y];
        
        n1.adjacent.add(n2);
        n2.adjacent.add(n1);
    }
    
    void addQuery(int x) {
        querys.add(x);
    }
    
    void makeSubTree(int root) {
        dp[root] = 1; // 서브 트리에 포함
        Node node = nodes[root]; // 노드 꺼낸 후
        
        for (Node ad : node.adjacent) {
            if (dp[ad.x] == 0) { // 처음 호출 되는 경우만 재귀적 실행  --> dp == 0
                makeSubTree(ad.x); // 인접한 노드 여기서는 자식 노드를 의미
                dp[root] += dp[ad.x]; // 자식 노드의 서브트리를 더함 
            }
        }
    }
    
    void getPrint() {
        StringBuilder sb = new StringBuilder();
        
        for (Integer q : querys) {
            sb.append(dp[q]).append("\n");
        }
        
        System.out.print(sb);
    }
    
}

 

2. 풀이 중점 사항

 

문제를 출제하신 분께서 힌트로 너무나 자세하게 잘 설명해 주셔서 추가로 설명드릴 부분은 makeSubTree를 구현한 메서드입니다.

 

void makeSubTree(int root) {
    dp[root] = 1; // 서브 트리에 포함
    Node node = nodes[root]; // 노드 꺼낸 후
    
    for (Node ad : node.adjacent) {
        if (dp[ad.x] == 0) { // 처음 호출 되는 경우만 재귀적 실행  --> dp == 0
            makeSubTree(ad.x); // 인접한 노드 여기서는 자식 노드를 의미
            dp[root] += dp[ad.x]; // 자식 노드의 서브트리를 더함
        }
    }
}

 

makeSubTree를 만드는 과정에서 각 노드는 서로 인접한 노드로 설정되어 있습니다. 따라서, 각 노드에 대한 visited를 설정해주어야 합니다. 만약 방문 여부를 설정하지 않는 다면 인접한 노드가 다시 root 노드를 호출하기 때문에 무한 루프에 빠지게 됩니다.

저는 이 과정에서 dp[i]의 값이 0이라면, 아직 재귀 호출이 되지 않은 상태 즉, 자식 노드이므로 dp[i] == 0인 경우를 visited로 대체하여 사용하였습니다.

 

dp를 이용하여 한 번에 서브 트리의 개수를 구할 수 있었습니다.

이상으로 백준 트리와 쿼리 자바 풀이를 마치도록 하겠습니다. 감사합니다!!