BOJ 25402(트리와 쿼리)풀이

이미지를 누르면 문제로 이동합니다

이 문제는 여러 가지 풀이 방법이 있지만 이 글에선 유니온 파인드(줄여서 유파)를 이용한 풀이를 사용하겠다.

 

이 문제를 간단히 요약하면 사용할 수 있는 정점이 주어지면 한 정점에서 다른 정점으로 갈 수 있는 모든 경우의 수를 구하라는 문제이다.

 

이때 n개의 정점들이 이어져있다고 해보자. n개의 정점들은 자신을 제외한 모든 정점에게 갈수 있다고 할 때 이 그룹에서 나올 수 있는 경우의 수는 n*(n-1)/2 개가 된다. 이 글에선 이걸 이용할 것이다.

 

일단 처음에 한 정점을 기준으로 번호를 다시 부여해준다.(이 글에선 1번 정점을 기준으로 했다.)

번호를 부여할때 자신의 부모 노드의 번호를 같이 저장한 뒤 입력을 받을 때 새로 부여된 자신의 번호와 부모의 번호를 함께 저장한다.

그 뒤 자신에게 부여된 번호를 오름차순으로 정렬하게 되면 번호를 부여할 때 탐색했던 순으로 정렬되게 된다.

 

만약 주어진 정점들 중에 자신의 부모 노드가 있고 그 부모 노드가 다른 정점과 연결되어있다면?

 

그러면 자식노드는 당연히 부모 노드와 연결이 될 것이다.

 

만약 아니라면 일정 번호까진 자신의 자식 노드들이 되게 되므로 현재 노드가 시작 노드가 되는 그룹이 새로 형성될 것이다.

이 과정을 입력받은 노드들에게 반복을 해주면 여러 개의 그룹이 형성될 것이다. 이때 각 그룹마다 위 식을 이용해 계산해준 뒤 더해주면 문제에서 원하는 답을 구할 수 있다.

 

 

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
struct nd{
    ll s,e;
};
ll a,b,c,cnt=0,q,k;
vector<ll>adj[250100];
vector<nd>dfsn(250100);
ll st[250100];
ll inp[250100];
ll parent[250100];
ll ck[250100];
ll ans[250100];
ll as=0;
void dn(ll v){
    cnt++;
    dfsn[v].s=cnt;
    for(int i=0;i<adj[v].size();i++){
        if(dfsn[adj[v][i]].s==0){
            st[cnt+1]=dfsn[v].s;
            dn(adj[v][i]);
        }
    }
    dfsn[v].e=cnt;
    
}
ll find(ll x){
    if(parent[x]<0){
        return x;
    }else{
        ll y=find(parent[x]);
        parent[x]=y;
        return y;
    }
}
void uni(ll x,ll y){
    x=find(x);
    y=find(y);
    if(x==y)return;
    if(parent[x]<parent[y]){
        parent[x]+=parent[y];
        parent[y]=x;
    }else{
        parent[y]+=parent[x];
        parent[x]=y;
    }
}
    
int main(){
    ios::sync_with_stdio(0);
    cin.tie(0),cout.tie(0);
    fill(parent,parent+250100,-1);
    cin >> a;
    ll idx=0;
    for(int i=0;i<a-1;i++){
        cin >> b >> c;
        adj[b].push_back(c);
        adj[c].push_back(b);
    }
    st[1]=-1;
    dn(1);
    cin >> q;
    for(int i=0;i<q;i++){
        as=0;
        idx=0;
        cin >> k;
        for(int g=0;g<k;g++){cin >> inp[g];inp[g]=dfsn[inp[g]].s;}
        sort(inp,inp+k);
        for(int g=0;g<k;g++){
            ck[inp[g]]=1;
            if(st[inp[g]]==-1){
                ans[idx]=inp[g];
                idx++;
            }else{
                if(ck[st[inp[g]]]==1){
                    uni(inp[g],st[inp[g]]);
                }else{
                    ans[idx]=inp[g];
                    idx++;
                }
            }
        }
        for(int g=0;g<idx;g++){
            ll uo=-parent[ans[g]];
            as+=uo*(uo-1)/2;
        }
        cout << as << "\n";     
        for(int g=0;g<k;g++){ck[inp[g]]=0;parent[inp[g]]=-1;}
    }
 
}