백준 15681 - 트리와 쿼리

문제 설명

간선에 가중치와 방향성이 없는 임의의 루트 있는 트리가 주어졌을 때, 아래의 쿼리에 답해보도록 하자.

• 정점 U를 루트로 하는 서브트리에 속한 정점의 수를 출력한다.

만약 이 문제를 해결하는 데에 어려움이 있다면, 하단의 힌트에 첨부한 문서를 참고하자.

https://www.acmicpc.net/problem/15681

 

 

---

 

 

 

 

 

 

제한 사항

---

 

 

 

 

 

 

풀이

문제를 요약하면, 주어진 정보로 트리를 만들고 쿼리로 주어진 노드의 서브트리의 개수를 구하는 것이다.

 

해당 문제에서는 트리를 만들 때, DFS를 통해 만들어주면 쉽게 문제를 풀 수 있다.

DFS를 통해 트리를 만들다 리프 노드에 도착하면 순회를 종료하고 부모로 올라간다.

 

예를 들어, 1번 노드는 리프 노드이기 때문에 해당 노드의 서브트리는 존재하지 않는다.

따라서, 본인만 카운팅한 1이 될 것이다.

즉, 리프 노드에 도착하면 본인을 1로 기록한 뒤 부모 노드로 거슬러 올라간다.

해당 과정을 반복하면 자신의 서브트리의 개수가 모두 기록된 배열을 하나 생성할 수 있다.

void DFS(vector<vector<int>>& edges, int node)
{
    SubTree[node] = 1;

    for (auto next : edges[node])
    {
        if (visited[next]) continue;
        visited[next] = true;

        DFS(edges, next);
        SubTree[node] += SubTree[next];
    }
}

 

이후, 쿼리의 내용을 처리하면 된다.

해당 과정은 미리 기록한 서브 트리의 개수를 읽어오면 되기 때문에 $O(1)$의 시간이 걸린다. 

for (int i = 0; i < Q; i++)
{
    int q;
    cin >> q;
    cout << SubTree[q] << "\n";
}

---

 

 

 

 

 

전체 코드

#include <stdio.h>
#include <cstring>
#include <string>
#include <vector>
#include <algorithm>
#include <iostream>
#include <sstream>
#include <cmath>
#include <climits>
#include <queue>
#include <map>
#include <unordered_map>
#include <set>

using namespace std;

int N, R, Q;

vector<int> SubTree;
vector<int> link;
vector<bool> visited;

void DFS(vector<vector<int>>& edges, int node)
{
    SubTree[node] = 1;

    for (auto next : edges[node])
    {
        if (visited[next]) continue;
        visited[next] = true;

        DFS(edges, next);
        SubTree[node] += SubTree[next];
    }
}

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

    cin >> N >> R >> Q;

    SubTree.resize(N + 1, 0);
    visited.resize(N + 1, false);
    vector<vector<int>> edges(N+1);

    for (int i = 1; i < N; i++)
    {
        int a, b;
        cin >> a >> b;

        edges[a].push_back(b);
        edges[b].push_back(a);
    }

    visited[R] = true;
    DFS(edges, R);

    for (int i = 0; i < Q; i++)
    {
        int q;
        cin >> q;
        cout << SubTree[q] << "\n";
    }

    return 0;
}

---