서브트리 크기

루트가 있는 트리에서 각 노드의 서브트리 크기(자신 포함)를 구한다.

---

 

 

#include <iostream>
#include <vector>

using namespace std;

vector<int> tree[100005];
bool visited[100005];
int sub_tree_size[100005];

int dfs(int cur) {
    visited[cur] = true;
    sub_tree_size[cur] = 1; 

    for (auto nxt : tree[cur]) {
        if (visited[nxt]) {
            continue;
        }
        sub_tree_size[cur] += dfs(nxt);
    }

    return sub_tree_size[cur];
}

int main() {
    ios_base::sync_with_stdio(false);                       
    cin.tie(NULL);

    int n;
    cin >> n;

    for (int i = 2; i <= n; i++) {
        int parent;
        cin >> parent;
        
        tree[parent].push_back(i);
        tree[i].push_back(parent);
    }

    dfs(1);

    for (int i = 1; i <= n; i++) {
        cout << sub_tree_size[i] << " ";
    }
    cout << "\n";

    return 0;
}

몇개의 노드를 입력 받을지 정한뒤 맨처음 1번 노드를 제외 한 노드들이 연결된 부모트리를 vector tree에 값을 초기화 한뒤 

dfs를 이용하여 검사한다 첫번째 자기 자신은 확인을 한 상태 이므로 검사를 한지 확인 하기 위해 bool형의 배열을 이용한다.

그 후 for문을 돌며 트리를 탐색 하는데 만약에 검사를 한것이면 continue를 하여 반복문을 나간다.

검사를 하지 않은 것 이라면 sub_tree_size에 값을 증가 시킨다. 

그후 dfs를 한번더 하여 값을 역으로 바꾼뒤 for문을 이용하여 출력 한다.