[ 백준 ] 1068 트리

📌 Problem

트리 구조에서 하나의 노드를 제거하였을 때(노드를 루트로 하는 서브 트리 전체가 삭제된다.) 남은 leaf node의 개수를 구한다.

📌 Solution

Node는 자신의 번호와 자식들의 번호를 가진다.

static class Node {
    int nno;
    List<Node> children = new ArrayList<>();
}

-1을 부모로 하는 node를 root node으로 가져온다. 나머지는 부모 node에 추가한다

Node root = null;
StringTokenizer tokenizer = new StringTokenizer(reader.readLine());
for (int nno = 0; nno < NODE_NUM; nno++) {
    int parentNno = Integer.parseInt(tokenizer.nextToken());
    if (parentNno == -1) root = nodes.get(nno);
    else nodes.get(parentNno).addChild(nodes.get(nno));
}
assert root != null;

Root node를 시작으로 BFS으로 탐색한다. 만약 탐색한 node가 제거된 node이면 넘어간다. 만약 탐색한 node가 자식을 가지고 있지 않거나 제거된 node 하나만을 자식으로 가지고 있으면 leaf node이다.

int removedNno = Integer.parseInt(reader.readLine());

int leafCountNum = 0;
Queue<Node> queue = new LinkedList<>(Collections.singleton(root));
while (!queue.isEmpty()) {
    Node node = queue.poll();
    if (node.nno == removedNno) continue;
    List<Node> children = node.children;
    if (children.size() == 0 || (children.size() == 1 && children.get(0).nno == removedNno)) 
        leafCountNum++;
    for (Node child : children) queue.offer(child);
}

📌 Code

// Step 1. Create nodes
final int NODE_NUM = Integer.parseInt(reader.readLine());
Map<Integer, Node> nodes = IntStream.range(0, NODE_NUM).boxed()
        .collect(Collectors.toMap(nno -> nno, Node::new));

// Step 2. Create tree
Node root = null;
StringTokenizer tokenizer = new StringTokenizer(reader.readLine());
for (int nno = 0; nno < NODE_NUM; nno++) {
    int parentNno = Integer.parseInt(tokenizer.nextToken());    // Node number of parent
    if (parentNno == -1) root = nodes.get(nno);                 // If parent is -1 then root
    else nodes.get(parentNno).addChild(nodes.get(nno));
}
assert root != null;

// Read removed node number
int removedNno = Integer.parseInt(reader.readLine());

// Step 3. Count leaf number by BFS
int leafCountNum = 0;
Queue<Node> queue = new LinkedList<>(Collections.singleton(root));
while (!queue.isEmpty()) {
    Node node = queue.poll();
    if (node.nno == removedNno) continue;   // If node is removed continue
    List<Node> children = node.children;

    // If node have no children or only removed node then current node is leaf node
    if (children.size() == 0 || (children.size() == 1 && children.get(0).nno == removedNno))
        leafCountNum++;

    // Add children to queue for bfs
    for (Node child : children) queue.offer(child);
}

result.append(leafCountNum);

static class Node {
    int nno;
    List<Node> children = new ArrayList<>();

    public Node(int nno) {
        this.nno = nno;
    }

    public void addChild(Node node) {
        children.add(node);
    }
}