백준 1068 - 트리

문제

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

리뷰

뭔가 처음 문제를 접하였을 때 유니온 파인드에서 부모를 찾는데 쓰는 로직을 응용하면
풀이가 쉬울 것 같다는 생각으로 접근을 하였다.

우선 각 노드의 부모를 저장하는 parents 배열과 각 노드의 자식의 수를 저장하는 childs
배열 두 개를 정의하였다. 부모 정보를 parents 배열에 저장한 후 setRemoveNodes
함수를 통하여 제거된 노드를 루트로 하는 서브트리에 속한 노드들의 childs 값을 모두 -1로
초기화해준다. 이후 countChilds 함수를 통해 childs 의 값이 -1이 아닌 노드만
자식 노드의 수를 카운팅해주었다. 이후 노드중 childs 의 값이 0인 경우만 카운팅하여
답을 도출했다.

로직의 시간복잡도는 최악의 경우에도 $O(N)$으로 수렴하므로 무난히 제한 조건을 통과할 수
있었다.

코드

import java.util.*;

import static java.lang.Integer.parseInt;


public class Main {
    static int[] parents;
    static int[] childs;

    public static void main(String[] args) {
        Scanner in = new Scanner(System.in);
        int N = parseInt(in.nextLine());

        parents = new int[N];
        childs = new int[N];

        StringTokenizer st = new StringTokenizer(in.nextLine());
        for (int i = 0; i < parents.length; i++)
            parents[i] = parseInt(st.nextToken());

        int removedNode = parseInt(in.nextLine());
        for (int node = 0; node < N; node++)
            setRemovedNodes(node, removedNode);

        for (int node = 0; node < N; node++) {
            if (childs[node] != -1)
                countChilds(node);
        }

        long count = Arrays.stream(childs)
                .filter(c -> c == 0)
                .count();

        System.out.println(count);
        in.close();
    }

    static void setRemovedNodes(int start, int removed) {
        int parent = start;

        while (parent != -1) {
            if (parent == removed) {
                childs[start] = -1;
                return;
            }

            parent = parents[parent];
        }
    }

    static void countChilds(int start) {
        int parent = parents[start];

        while (parent != -1) {
            childs[parent]++;

            parent = parents[parent];
        }
    }
}

결과