[BOJ / C++] 트리 (트리)

문제링크: https://www.acmicpc.net/problem/1068

 

1. 코드

#include <bits/stdc++.h>
using namespace std;

int main() {
    ios::sync_with_stdio(0); cin.tie(0);

    int n;
    cin >> n;

    vector<int> parent(n);
    vector<vector<int>> adj(n+1);
    int root;
    for (int i = 0; i < n; i++) {
        cin >> parent[i];

        if (parent[i] == -1) {
            root = i;
        }
        else {
            adj[parent[i]].push_back(i);
        }
    }

    int deleteNode;
    cin >> deleteNode;

    // 끊어내기
    if (deleteNode == root) {
        cout << 0;
        return 0;
    }
    auto it = find(adj[parent[deleteNode]].begin(), adj[parent[deleteNode]].end(), deleteNode);
    adj[parent[deleteNode]].erase(it);

    // leaf count
    int leaf = 0;
    queue<int> Q;
    Q.push(root);

    while (!Q.empty()) {
        int cur = Q.front(); Q.pop();

        if (adj[cur].size() == 0) leaf++;

        for (int nxt : adj[cur]) {
            Q.push(nxt);
        }
    }

    cout << leaf;
}

 

 

2. 시간복잡도

\(n=50\)이고 인접리스트에서 BFS 탐색을 하므로 시간복잡도는 \(O(v+e) = 50 + 49 < 10^{8}\)이므로 시간제한 2초 내에 통과 가능하다.

 

 

3. 분석

root를 따로 저장하고 adj list에는 부모에서 자식으로 가는 edge만 남겨놓는다. 

이후 deleteNode를 입력받고 해당 부분을 끊어내는데 두가지 방법이 가능하다.

 

(1) 직접 adj list에서 없애기

실제 위 코드대로 adj에서 찾아내서 삭제한다. 루트의 경우 -1이라 OOB가 나므로 예외처리한다.

// 끊어내기
if (deleteNode == root) {
    cout << 0;
    return 0;
}
auto it = find(adj[parent[deleteNode]].begin(), adj[parent[deleteNode]].end(), deleteNode);
adj[parent[deleteNode]].erase(it);

 

(2) BFS에서 deleteNode를 `continue` 하고 탐색하기

if (cur == deleteNode) continue;

 

BFS에서 해당 노드를 만났을 때 더이상 BFS 탐색을 하지 않는 방식으로도 해결 가능하다. 이 방법이 조금 구현은 간단하다.