백준 15681 (트리와 쿼리) c++

문제

15681번: 트리와 쿼리

최상단 루트 노드가 주어지고 트리의 연결 정보가 주어졌을 때 노드의 번호를 입력하면 입력된 노드 번호를 포함한 

subtree의 갯수를 출력하는 문제입니다.

예를 들어 루트 노드가 5이고 입력된 노드가 5, 4, 8 이라면

답은 각각 9, 4, 1 이 됩니다.

---

먼저 트리의 정보를 백터의 인접리스트 형태로 저장해줍니다.

for (int i = 1; i < n; i++) {
        cin >> a >> b;
        v[a].push_back(b);
        v[b].push_back(a);
    }

 

이제 각 노드마다 subtree의 갯수 정보를 dp배열에 저장해 줄것입니다.

예를 들어 dp[n] = i 라고 하면 n번째 노드는 i개의 subtree갯수를 가지고 있다는 의미입니다.

우리는 각 노드 마다 subtree갯수 정보를 저장하기 위해서 dfs순회 방식을 사용할 것입니다.

void dfs(ll x, ll par) {
    dp[x] = 1; //자기 자신은 항상 포함하기 때문
    for (auto nxt : v[x]) {
        if (nxt == par)continue;
        dfs(nxt, x);
        dp[x] += dp[nxt];
    }
}

main함수에서는 dfs(r, 0) 이 실행되었습니다. 

즉, 첫번째 인자에는 루트노드가 들어가고 두번째 인자에는 부모 노드가 들어가게 됩니다.

우리가 처음에 전체 트리의 루트 노드는 5라고 하였으므로 r에는 5가 들어갈 것이고 

5는 전체트리의 루트노드이기 때문에 부모노드가 없으므로 그냥 0이 들어가게 됩니다.

 

2번째 줄에서 보시다 시피 모든 노드들의 subtree는 자기자신을 포함하므로 1로 초기화 해줍니다.

3번째 for문을 보면 auto nxt : v[x] 라는 의미는 v[x]의 모든 원소들을 반복하겠다는 의미입니다.

4번째 줄에서 nxt==par이면 continue입니다. 우리는 현재 자식노드만 탐색하면 되기 때문입니다.

---

전체 코드 입니다.

#include<iostream>
#include<algorithm>
#include<climits>
#include<string>
#include<vector>
#include<queue>
#include<stack>
#include<deque>
#include<cmath>
#include<time.h>
#include<cstring>
#include<cmath>
#include<cstring>
using namespace std;
using ll = long long;
using ull = unsigned long long;
ll dp[100001];
vector<ll>v[100001];
ll n, r, q;
ll a, b;
void dfs(ll x, ll par);
int main() {
    ios_base::sync_with_stdio(false), cin.tie(NULL), cout.tie(NULL);
 
    cin >> n >> r >> q;
    for (int i = 1; i < n; i++) {
        cin >> a >> b;
        v[a].push_back(b);
        v[b].push_back(a);
    }
    dfs(r, 0);
    while (q--) {
        cin >> a;
        cout << dp[a] << '\n';
    }
 
 
    return 0;
}
void dfs(ll x, ll par) {
    dp[x] = 1; //자기 자신은 항상 포함하기 때문
    for (auto nxt : v[x]) {
        if (nxt == par)continue;
        dfs(nxt, x);
        dp[x] += dp[nxt];
    }
}