백준 1068번 트리 java

오늘은 2일전에 풀었던 트리라는 문제를 다시 한번 풀어보았다. 이 문제는 설명이 친절하지 않아서 많은 혼란을 주었던 문제이다.

🔥 문제

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

트리에서 리프 노드란, 자식의 개수가 0인 노드를 말한다.

트리가 주어졌을 때, 노드 하나를 지울 것이다. 그 때, 남은 트리에서 리프 노드의 개수를 구하는 프로그램을 작성하시오. 노드를 지우면 그 노드와 노드의 모든 자손이 트리에서 제거된다.

예를 들어, 다음과 같은 트리가 있다고 하자.

현재 리프 노드의 개수는 3개이다. (초록색 색칠된 노드) 이때, 1번을 지우면, 다음과 같이 변한다. 검정색으로 색칠된 노드가 트리에서 제거된 노드이다.

이제 리프 노드의 개수는 1개이다.

입력

첫째 줄에 트리의 노드의 개수 N이 주어진다. N은 50보다 작거나 같은 자연수이다. 둘째 줄에는 0번 노드부터 N-1번 노드까지, 각 노드의 부모가 주어진다. 만약 부모가 없다면 (루트) -1이 주어진다. 셋째 줄에는 지울 노드의 번호가 주어진다.

출력

첫째 줄에 입력으로 주어진 트리에서 입력으로 주어진 노드를 지웠을 때, 리프 노드의 개수를 출력한다.

문제 풀이

이 문제는 dfs 탐색을 통해 해결하였다.

우선 그래프를 생성하고 부모 자식 간에 연결을 하였다. 단, 루트 노드일 경우는 따로 루트 노드를 저장하였다. 처음에는 루트가 무조건 0인 줄 알고 따로 저장하지 않고 dfs 함수에 0을 하드코딩으로 넣었는데 이 점을 주의해야 한다.

그리고 방문 리스트는 따로 생성하지 않았는데 이유는 트리여서 단방향 dfs이므로 굳이 방문 체크하지 않아도 모두 탐색할 것이기 때문이다. 이렇게 그래프를 생성하고 연결한 후 제거할 루트 노드 값을 받고 제거를 하였다.

제거할 때는 한곳만 끊어버리면 나머지 아래쪽은 탐색을 하지 않을 것이기 때문에 굳이 자식에 자식까지 처리를 해주지 않았다.

그래서 최종적으로 아래와 같은 코드로 문제를 해결하였다.

코드

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

class Main {

    static int N;
    static ArrayList<ArrayList<Integer>> graph = new ArrayList<>();
    static boolean[] visited;

    // 루트 노드
    static int root = 0;

    // 제거해야할 노드
    static int removeNode;
    static int result = 0;
    public static void main(String[] args) throws IOException {

        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        N = Integer.parseInt(br.readLine());
        visited = new boolean[N];

        for (int i = 0; i < N; i++) {
            graph.add(new ArrayList<>());
        }

        // 그래프 연결
        StringTokenizer edge = new StringTokenizer(br.readLine());
        for (int i = 0; i < N; i++) {
            int parent = Integer.parseInt(edge.nextToken());
            if (parent != -1) {
                graph.get(parent).add(i);
            } else {
                root = i;
            }
        }

        removeNode = Integer.parseInt(br.readLine());
        if (removeNode == root) {
            System.out.println(0);
            return;
        }

        // 노드 제거
        for (int i = 0; i < N; i++) {
            graph.get(i).remove(Integer.valueOf(removeNode));
        }

        dfs(root);
        System.out.println(result);
    }

    private static void dfs(int node) {
    
    	// 자식이 없으면 리프노드이므로 결과에 +1처리 후 return
        if (graph.get(node).isEmpty()) {
            result++;
            return;
        }

        for (int next : graph.get(node)) {
            dfs(next);
        }
    }
}