트리와 쿼리( BOJ 25402 )

문제 : https://www.acmicpc.net/problem/25402

25402번: 트리와 쿼리

 

문제 파악하기

집합 S에 포함된 정점들이 몇 개씩 연결되었는지 판단하는 문제입니다. 문제에서 요구하는 연결 강도는 결국 연결된 정점 중 2개의 정점을 고를 수 있는 경우의 수를 의미하며, 만약 t개의 정점이 연결되어 있다면 연결 강도는 t*(t-1)/2 로 계산됩니다. 따라서, 문제를 해결하기 위해서는 몇 개의 정점들이 현재 연결되어 있는지 확인해야 합니다. 다만, 모든 경우를 모두 계산하기에는 최대 250,000개의 정점과 최대 1,000,000개의 질문이 너무 많기에 효율적인 알고리즘이 필요합니다. 이 문제에서는 다행히 그래프가 트리 구조이기 때문에 트리 구조의 특징을 활용할 수 있습니다.

 

문제 해결하기

트리는 항상 정점보다 하나 적은 간선의 개수를 가지고 있으며, 모든 정점 사이의 경로는 딱 1가지 있습니다. 그리고 이 경로는 항상 [ 부모 정점 - 자식 정점 ] 으로 연결되어 있습니다. 따라서, 우리는 주어진 그래프를 트리로 만든 다음, 집합 S의 정점 중 부모 정점이 S에 포함된 정점들만 서로 연결(Union)해주면 됩니다. 여기서 연결(Union)한다는 의미는 분리 집합 연산 중 결합(Union) 연산을 의미하며, 연결 시 각 그룹의 노드 개수 역시 갱신해주면 됩니다.

 

질문이 하나인 경우에는 간단한게 모든 값을 초기화하여 계산할 수 있지만, 이 문제의 경우에는 정점의 수도 많고, 질문의 수도 많습니다. 그렇기에 최대한 초기화 하는 시간을 줄여야 제한 시간 이내에 정답을 구할 수 있습니다. 그렇기에 가장 효율적으로 초기화하는 방법은 집합 S에서 사용되는 정점만 초기화 하면 됩니다. 다른 정점들은 아에 방문을 하지 않기에 따로 초기화 해줄 필요가 없기 때문입니다. 자세한 소스코드는 다음과 같습니다.

 

소스코드

#include <stdio.h>
#include <string.h>
#include <vector>
#define NMAX 250010
#define lint long long int
using namespace std;
 
int N, Q, a, b;
int M, K, t;
int inp[NMAX];
 
int treeParent[NMAX];
vector< int > graph[NMAX], tree[NMAX];
 
int p, pp;
int used[NMAX], check[NMAX], parent[NMAX];
 
lint ret;
lint cnt[NMAX];
 
int find(int p) {
    if(parent[p] == p) return p;
    else return parent[p] = find(parent[p]);
}
 
void make_tree(int idx) {
    check[idx] = 0;
    for(int child:graph[idx]) {
        if(!check[child]) continue;
        
        tree[idx].push_back(child);
        make_tree(child);
        treeParent[child] = idx;
    }
}
 
int main() {
    // input
    scanf("%d", &N);
    for(int i=1;i<N;i++) {
        scanf("%d %d", &a, &b);
        graph[a].push_back(b);
        graph[b].push_back(a);
    }
    
    // make tree
    memset(check, -1, sizeof(check));
    make_tree(1);
    
    scanf("%d", &M);
    for(int i=1;i<=M;i++) {
        // init
        ret = 0;
        
        // query
        scanf("%d", &K);
        for(int j=0;j<K;j++) {
            scanf("%d", &inp[j]);
            
            // init
            used[inp[j]] = i;
            cnt[inp[j]] = 1;
            parent[inp[j]] = inp[j];
        }
        
        // UF(자식과 부모 모두 S에 포함된 경우만 합치기)
        for(int j=0;j<K;j++) {
            if(used[inp[j]] == i and used[treeParent[inp[j]]] == i) {
                p = find(inp[j]);
                pp = find(treeParent[inp[j]]);
                
                cnt[pp] += cnt[p];
                parent[p] = pp;
            }
        }
        
        // check
        for(int j=0;j<K;j++) {
            if(parent[inp[j]] == inp[j]) {
                ret += cnt[inp[j]]*(cnt[inp[j]]-1)/2;
            }
        }
        
        // print
        printf("%lld\n", ret);
    }
}

후기

단순한 분리 집합 문제가 아닌, 트리의 특성을 활용한 문제입니다. 문제 자체는 단순하지만 초기화 하는 방법이나 트리의 특성을 생각할 거리가 있는 문제이기에 재미있는 문제라고 생각합니다. 트리를 배운 사람에게 추천할 만한 문제라고 생각합니다.