[C++] 백준 온라인 저지 1068번 트리 풀이

트리의 리프 노드가 몇 개인지 찾는 문제

 

---

문제

 

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

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

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

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

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

 

입력

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

 

출력

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

 

 

풀이

 

 이 문제의 주의사항은 입력으로 주어지는 부모 노드의 순서가 꼭 루트 노드부터 주어지지 않는 것입니다. 예제는 모두 -1로 시작하니까 당연히 맨 처음에 나오는 노드가 루트일거라 생각하고 입력 자체를 잘못 받았더니 예제는 맞았지만 제출하니까 시작하자마자 틀렸습니다 ㅠ.ㅠ 예제 입력에 이런 거 보완 좀... 

 루트 노드가 입력 중간에 나올 수도 있기 때문에 입력 중간에 -1이 나오는 i번째 노드를 루트 노드로 지정해 주어야 합니다. 이거 때문에 의미 없는 제출을 몇 번을 한 건지...

 

 그리고 자식 노드를 하나만 가지고 있는데 지워야 할 노드가 그 자식 노드인 경우에 부모 노드를 리프 노드로 카운트 해 주어야 합니다. 

 

 그 외에는 단방향 그래프로 구성한 다음 DFS로 탐색하면서 리프 노드라면 카운트하고 지워야 할 노드라면 방문하지 않는(지워야 할 노드를 지우기 위한 다른 처리를 할 필요가 없음) 코드를 짜면 통과할 수 있습니다. 

 

 

코드

 

#include <iostream>
#include <vector>

using namespace std;

int N;
int nodeToRemove;
int answer = 0;
vector<vector<int>> tree;
vector<bool> visited;

void dfs(int cur)
{
    //자식이 없으면 리프 노드
    if (0 == tree[cur].size())
    {
        answer++;
        return;
    }
    
    //자식이 하나인데 지워야 할 노드라면 현재 노드가 리프 노드가 됨
    if (1 == tree[cur].size() && nodeToRemove == tree[cur][0])
    {
        answer++;
        return;
    }
    
    for (int i = 0; tree[cur].size() > i; i++)
    {
        auto next = tree[cur][i];
        //다음 노드를 방문하지 않았으면서 지워야 할 노드가 아닐 때 방문
        if (!visited[next] && nodeToRemove != next)
        {
            visited[next] = true;
            dfs(next);
        }
    }
}

int main()
{
    cin >> N;
    
    int root;
    tree.assign(N, vector<int>(0));
    for (int i = 0; N > i; i++)
    {
        int input;
        cin >> input;
        
        //입력이 -1이라면 루트 노드
        if (-1 == input) root = i;
        else
            tree[input].emplace_back(i);
    }
    
    cin >> nodeToRemove;
    
    if (nodeToRemove != root)
    {
        visited.assign(N, false);
        visited[root] = true;
        dfs(root);
    }
    
    cout << answer;
    
    return 0;
}