[15681] 트리와 쿼리

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

15681 - 트리와 쿼리

본문

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

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

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

입력

트리의 정점의 수 $N$과 루트의 번호 $R$, 쿼리의 수 $Q$가 주어진다. ($2 \leq N \leq 10^5$, $1 \leq R \leq N$, $1 \leq Q \leq 10^5$)

이어 $N-1$줄에 걸쳐, $U$ $V$의 형태로 트리에 속한 간선의 정보가 주어진다. ($1 \leq U$, $V \leq N$, $U \neq V$)

이는 $U$와 $V$를 양 끝점으로 하는 간선이 트리에 속함을 의미한다.

이어 $Q$줄에 걸쳐, 문제에 설명한 $U$가 하나씩 주어진다. ($1 \leq U \leq N$)

입력으로 주어지는 트리는 항상 올바른 트리임이 보장된다.

출력

Q줄에 걸쳐 각 쿼리의 답을 정수 하나로 출력한다.

제한

시간 제한	메모리 제한
1sec	128MB

• $1 \leq N,Q \leq 10^5$

---

풀이

노드가 10만개, 쿼리도 10만개. 트리가 1자 모양으로 만들어진다면 그때그때마다 쿼리를 1초 안에 해결할 수는 없다. 결국 트리를 만들면서 쿼리마다의 답을 미리 기록하는 DP 방법을 사용해야한다는 점을 알 수 있다.

트리의 간선을 모두 입력으로 주고, 트리의 root마저 문제에서 제공한다. 심지어 트리의 모양은 항상 올바르다는 정보까지 보장한다. 간선들만 이어주면 트리 모양임은 보장하고, root 또한 정해져있기 때문에 부모 / 자식 관계가 확실히 나누어져있음을 알 수 있다. 본인과 연결된 노드들 중 부모 1개를 제외한 나머지 노드 전부가 본인의 자식 노드가 되는 것이다. 이것은 직접 예제를 그려보며 이해하자.

트리 구조를 만드는 방법은 편법을 이용하자. 어짜피 트리 구조는 완전함을 보장하므로, 간선들을 이어 각각의 관계를 만들지 말고 그 간선들만을 기록해두는 것이다. 이후에 현재 노드를 탐색할 때, 간선이 있다면 연결된 다른 노드는 부모 또는 자식 노드임을 알 수 있다. 부모 노드는 이전 탐색에서 전달받는 노드 1개이므로, 나머지는 모두 자식 노드로 생각할 수 있다.

문제에서 원하는 답인 현재 노드의 서브트리의 정점의 개수를 반환하는 함수를 하나 만들도록 하자. 함수로 만드는 이유는, 이것은 재귀적으로 반복하면서 dp배열에 저장해두기 위함이다. 값을 기록해둘수도 있고, 트리 구조이기 때문에 dfs 방식으로 탐색을 진행하는 것이 효율적이기 때문이다. 서브트리의 정점의 개수는 $1(\text{본인}) + \sum \text{자식 노드의 서브트리 정점의 수}$로 계산할 수 있다.

• 참고 알고리즘 : 다이나믹 프로그래밍

코드

사용 언어 : C

최종 수정일 : 2023 / 1 / 22

#include <stdio.h>
#include <stdlib.h>

typedef struct Node {
    int x;
    struct Node* next;
} ND;

#define MAX_IDX (int)(1e5 + 1)
#define NONE 0

ND* tree[MAX_IDX];
int dp[MAX_IDX];
int n, root, q;

ND* createNewNode(int x) {
    ND* newNode = (ND*)malloc(sizeof(ND));
    newNode->x = x;
    newNode->next = NULL;
    return newNode;
}

void makePath(int a, int b) {
    ND* newNode = createNewNode(b);
    newNode->next = tree[a];
    tree[a] = newNode;
    return;
}

int getDP(int x, int parent) {
    ND* cur = tree[x];
    if (cur == NULL) {
        return (dp[x] = 1);
    }

    int retval = 1;
    while (cur != NULL) {
        if (cur->x != parent) {
            retval += getDP(cur->x, x);
        }
        cur = cur->next;
    }
    return dp[x] = retval;
}

int main() {
    scanf("%d %d %d", &n, &root, &q);
    for (int i = 1; i < n; ++i) {
        int a, b;
        scanf("%d %d", &a, &b);
        makePath(a, b);
        makePath(b, a);
    }

    getDP(root, NONE);

    for (int i = 0; i < q; ++i) {
        int a;
        scanf("%d", &a);
        printf("%d\n", dp[a]);
    }
    return 0;
}

• tree[] : 트리 구조의 간선을 인접리스트로 기록
• dp[] : 쿼리의 답을 미리 기록하는 dp배열