[BOJ] 1068 : 트리

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

1068번: 트리

 

문제 유형 : DFS / 트리

solved.ac 난이도 : Gold V

 

문제

 

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

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

 

입력

 

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

 

출력

 

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

 

 

내가 생각해본 플랜

 

둘째 줄에 각 노드의 부모 정보가 들어있으므로, 역으로 부모에서 어느 자식으로 이어지는지에 대한 인접 리스트를 만들어준 다음,

제거할 노드의 연결 정보를 모두 없애서 서브트리까지 모두 없애고, 부모 노드에서도 자식 노드르 제거할 노드를 없앤다.

그 뒤, dfs를 통해서 리프노드의 개수를 구한다.

 

 

막혔던 부분과 해결

 

처음으로 만든 코드를 제출하자, 단 1초도 걸리지 않아 틀렸다고 나온다. 문제에서 기본 제공하는 테스트 케이스는 루트가 항상 0번이었지만, 이는 함정이다. 루트가 꼭 0번이라는 법이 없었다. 이를 알아채고, 루트를 구하는 함수를 만들고, 탐색을 0번부터가 아닌, 루트부터 시작시켜서 문제를 해결하였다.

 

 

코드(Python)

import sys
input = sys.stdin.readline
#입력값 받기와 기본 그래프 생성
n = int(input())
nodes = list(map(int, input().split()))
del_node = int(input())
graph = [[] for _ in range(n)]
#노드의 부모 정보들로 인접리스트 만들기
for i in range(n):
    if nodes[i] != -1:
        graph[nodes[i]].append(i)
#그래프에서 노드 제거하기
def delete(node):
    for i in graph[node]:
        delete(i)
    graph[node] = []
    for i in range(n):
        if node in graph[i]:
            graph[i].remove(node)

delete(del_node)

# 루트 노드를 찾는 함수
def find_root(nodes):
    return nodes.index(-1)

# 루트 노드부터 탐색하여 리프 노드의 개수를 계산
def count_leaf(graph, v):
    if len(graph[v]) == 0:
        return 1
    cnt = 0
    for i in graph[v]:
        cnt += count_leaf(graph, i)
    return cnt

#루트를 찾고, 루트를 없애는 문제라면 바로 0을 반환하게 함
root = find_root(nodes)
print(count_leaf(graph, root) if del_node != root else 0)

 

코드(C++)

#include <iostream>
#include <vector>

using namespace std;

int n, del_node;
vector<int> nodes, graph[50];

void delete_node(int node) {
    for(int i = 0; i < graph[node].size(); i++){
        int child = graph[node][i];
        delete_node(child);
    }
    graph[node].clear();
    for (int i = 0; i < n; i++){
        for (int j = 0; j < graph[i].size(); j++){
            if(graph[i][j] == node){
                graph[i].erase(graph[i].begin() + j);
                break;
            }
        }
    }
}

int find_root(const vector<int> &nodes){
    for (int i = 0; i < nodes.size(); i++){
        if(nodes[i] == -1){
            return i;
        }
    }
    return -1;
}

int count_leaf(const vector<int> graph[], int v){
    if(graph[v].empty()){
        return 1;
    }
    int cnt = 0;
    for(int i = 0; i < graph[v].size(); i++){
        cnt += count_leaf(graph, graph[v][i]);
    }
    return cnt;
}

int main(){
    cin >> n;
    nodes.resize(n);
    for (int i = 0; i < n; i++){
        cin >> nodes[i];
        if(nodes[i] != -1){
            graph[nodes[i]].push_back(i);
        }
    }
    cin >> del_node;
    delete_node(del_node);
    int root = find_root(nodes);
    if(del_node == root){
        cout << 0 << endl;
    }
    else{
        cout << count_leaf(graph, root) << endl;
    }
    return 0;
}