[DFS] 백준 - 1068번: 트리 Java 풀이

[공부용이기 때문에 코드가 깔끔하지 않을 수 있습니다!]

 

1068번: 트리

 

 

문제는 트리에서 노드 하나를 삭제하고 그 트리의 리프 노드를 구하는 것이다.

그냥 마지막으로 입력받은 노드를 부모 노드에서 삭제시키고 dfs 돌리면 된다.

따라서 쉬운 편인 문제이다.

 

코드

import java.util.ArrayList;
import java.util.Scanner;

public class Main {
    static ArrayList<Integer>[] map;
    static boolean[] visit;
    static int n;
    static int count;

    public static void main(String[] args) {
        Scanner sc = new Scanner(System.in);

        n = sc.nextInt();

        map = new ArrayList[n];
        visit = new boolean[n];
        int root = 0;
        count = 0;

        for (int i = 0; i < n; i++) {
            map[i] = new ArrayList<>();
        }

        for (int i = 0; i < n; i++) {
            int a = sc.nextInt();

            if (a == -1)
                root = i;
            else
                map[a].add(i);
        }

        int d = sc.nextInt();

        if (d == root) {
            System.out.println(0);
            return;
        }

        for (int i = 0; i < n; i++) {
            if (map[i].contains(d)) {
                map[i].remove(Integer.valueOf(d));
            }
        }

        dfs(root);

        System.out.println(count);
    }

    private static void dfs(int start) {
        if (map[start].size() == 0) {
            count++;
            return;
        }

        visit[start] = true;

        for (int i = 0; i < map[start].size(); i++) {
            if (!visit[map[start].get(i)]) {
                dfs(map[start].get(i));
            }
        }
    }
}

 

위에서 말했던 것처럼 삭제할 노드의 부모에서 그 노드를 삭제 시키고 dfs를 돌면서 자식 노드가 없을 경우 즉, 리프 노드일 경우 카운팅을 해주고 dfs를 탈출한다.

 

그리고 입력받은 값이 만약 루트 노드와 동일하다면 그냥 0 출력하고 return 해주면 된다.