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

15681번: 트리와 쿼리

dfs와 dp를 이용한 문제이다. 입력으로 U와 V의 연결을 알려주는데 어느 방향이 루트로 향하는 방향인지 알 수가 없다. 그렇기에 양방향으로 입력을 받은 후 createTree를 통해 루트에서 자식 노드로 다시 방향을 구해주었다. 그리고 서브트리의 크기를 구하기 위해 루트에서 시작해 자식 노드를 경유하면서 서브 트리의 크기를 구해주었다. 그 후 크기들을 출력해주었다. 처음에는 사이즈를 dp방식으로 구해주지 않고 반복문을 돌며 카운트를 해주는 방식으로 구해주었더니 시간 초과가 났었다. 그래서 dp를 이용해 사이즈를 구해보는 방식으로 변경을 한 후 통과를 할 수 있었다. '먼저 완전 탐색을 해본 후 시간 초과가 난다면 dp를 이용해보자'고 생각했던 방식대로 문제를 해결하여서 뿌듯했다. 이러한 방식을 잘 인지해두어야 겠다.

#include <iostream>
#include <vector>
#include <cstring>

using namespace std;

int N, R, Q, result;
vector<int> n[100001];
vector<int> q;
vector<int> tree[100001];
int treeSize[100001];

void createTree(int parent, int now) {
    for (int i = 0; i < n[now].size(); i++) {
        int next = n[now][i];

        if (parent == next) continue;

        tree[now].push_back(next);
        createTree(now, next);
    }
}

void findSize(int node) {
    treeSize[node] = 1;

    for (int i = 0; i < tree[node].size(); i++) {
        int next = tree[node][i];

        findSize(next);
        treeSize[node] += treeSize[next];
    }
}

void solution() {
    createTree(0, R);
    findSize(R);

    for (int i = 0; i < q.size(); i++) {
        cout << treeSize[q[i]] << "\n";
    }
}

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

    cin >> N >> R >> Q;

    int U, V;
    for (int i = 0; i < N - 1; i++) {
        cin >> U >> V;
        n[U].push_back(V);
        n[V].push_back(U);
    }

    int a;
    for (int i = 0; i < Q; i++) {
        cin >> a;
        q.push_back(a);
    }

    solution();

    return 0;
}