자바랑 코덕코덕

[문제]

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

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

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

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

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

[문제 풀이]

제거할 노드를 부모로 가지고 있는 모든 노드들을 제거하기 위해 BFS를 사용했다.

delete(int node) 함수는 bfs를 구현한 것으로 (제거할 노드 == 부모 노드) 일때 큐에 넣어주고 하나씩 꺼내면서 제거해 주었다.

제거하면서 boolean 변수 visited=true로 방문체크를 하였다.

이후 leaf 노드를 찾기 위해 방문을 안한 노드를 찾아가면서 노드가 가리키는 숫자는 부모노드이므로 그 숫자들을 방문 체크하였다.

결국 하나도 방문하지 않은 노드가 leaf 노드가 된다.

방문하지 않는 노드 수를 세서 출력해주면 leaf 노드의 수가 나온다.

[코드]

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.Arrays;
import java.util.LinkedList;
import java.util.Queue;
import java.util.StringTokenizer;

public class Main {

    static int N;
    static int[] parents;
    static int removeNode;
    static boolean[] visited;
    static boolean[] leaves;

    static int stoi(String s) {
        return Integer.parseInt(s);
    }

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        N = stoi(br.readLine());

        parents = new int[N];
        visited = new boolean[N];
        leaves  = new boolean[N];
        Arrays.fill(leaves, true);

        StringTokenizer st = new StringTokenizer(br.readLine());
        for (int i=0; i<N; i++) {
            parents[i] = stoi(st.nextToken());
        }

        removeNode = stoi(br.readLine());
        if (parents[removeNode] == -1) {
            System.out.println(0);
            return;
        }

        delete(removeNode);

        for (int i=0; i<N; i++) {
            if (parents[i] == -1) continue;
            if (!visited[i]) {
                leaves[parents[i]] = false;
            }
        }
        
        int ans = 0;
        for (int i=0; i<N; i++) {
            if (leaves[i]) ans++;
        }
        System.out.println(ans);
    }

    static void delete(int node) {
        visited[node] = true;
        leaves[node] = false;

        Queue<Integer> queue = new LinkedList<>();
        queue.add(node);

        while(!queue.isEmpty()) {
            int cur = queue.poll();
            for (int i=0; i<N; i++) {
                if (cur == parents[i] && !visited[i]) {
                    visited[i] = true;
                    leaves[i] = false;
                    queue.offer(i);
                }
            }
        }
    }
}

[실행 결과]

실행 결과는 위와 같다.