BOJ 18473 - Fast Spanning Tree

※ 이 문제를 풀기 전에는 이 문제 를 미리 풀고 오시는 것을 추천드립니다. 

※ 이 포스팅을 보시기 전에는 이 포스팅 을 미리 보고 오시는 것을 추천드립니다. 

※ 이 포스팅은 BOJ 14435 놀이기구 2의 풀이를 모두 알고 있다고 가정하고 작성되어 있습니다.

문제 링크 : https://www.acmicpc.net/problem/18473 (R5)

문제 태그

union & find, smaller to larger

문제 풀이

일단 가장 간단하게 시간복잡도에 상관없이 문제를 푸는 방법에 대해서 생각을 해보자. 1부터 m번까지의 간선을 모두 보고 되는 간선을 체크하자.

당연하게도 우리는 가능한 간선들이 바뀔 때마다 계속해서 번호가 가장 간선을 선택해야하므로, 수를 추가하거나 빼면서 이를 관리하고, 그 중에서 가장 작은 수를 계속 뽑을 수 있는 자료구조인 우선 순위 큐를 사용하여 조건을 만족한, 가능한 간선들의 번호를 관리할 것이다. 

우선 순위 큐에서 번호가 가장 낮은 간선을 꺼내고, 간선의 양 끝점이 이미 같은 컴퍼넌트에 속해있는지 확인해야한다. 당연히 이 과정은 spanning tree를 만들 때 처럼 union&find를 사용해 줄 것이다. 

어떤 간선을 사용하겠다고 확정을 지으면, 추후에 답을 출력하기 위해서 큐에 넣어준다. 또한 이 연결로 인해서 어떤 노드가 속해있는 컴퍼넌트의 가중치의 합이 바뀌게 된다. 따라서 m개의 간선 중 추가로 조건을 만족한 간선이 있는지를 모두 확인해 주면 되고, 위와 같은 과정을 계속 반복하면 된다. 

이 경우 시간복잡도는 $O(MN)$이 걸리게 된다.

여기서 놀이기구 2 처럼 한가지 모토를 가지고 해보자. "확인할 간선만 확인한다." 즉, 간선을 영향을 받는 정점 두개에 대해서만 관리한다. 또한 어떤 값을 충족하기 위해서 적어도 둘 중 하나는 남은 값의 절반이상 만큼 증가해야한다.라는 관찰을 통해서 똑같이 일정 갯수만큼의 간선만 확인할 수 있다. 

하지만 여기서 생기는 문제점이 있다. 놀이기구 2의 상황을 생각해보면 매일마다 값이 바뀌는 요소는 딱 한개이다. 하지만 당연하게도 이 문제는 훨씬 난이도가 높은 문제이다. 한 번 간선을 이을 때마다 (업데이트 될 때마다) 최대 N개의 값이 변경 될 수 있다. 즉, N개의 간선 저장소를 모두 확인해야 한다는 것이다. 따라서 각 간선은 최대 lgS번 확인이 되어 이 시간 복잡도는 $O(MlgS)$이지만, 업데이트 할때마다 N개의 저장소를 확인하므로 이는 $O(N^2)$의 시간복잡도를 가지게 된다. 

여기서 아이디어를 내보자. 나는 업데이트가 될 때마다 N개의 모든 간선 저장소를 방문하기 싫다. 한가지 더 생각을 해보자. 당연하게도 "컴퍼넌트의 가중치의 합"을 가지고 비교하는 것이므로, 같은 컴퍼넌트 내에서는 같은 값을 가지게 된다. 즉 a-c 를 잇는 간선을 확인하기 위해서 a,c를 방문하는 대신 a와 b가 같은 컴퍼넌트에 있다면 b-c를 잇는다고 생각하고 b와 c를 방문하면 된다. 즉 같은 컴퍼넌트 내에서는 간선 저장소를 공유시킬 수 있다는 것이다.

당연히 간선저장소는 그룹의 최고 조상이 담당을 하고 있을 것이다. 컴퍼넌트내의 임의의 원소에서도 최고 조상에 무조건 도달할 수 있기 때문이다. 그러면 마지막으로 생각을 해야할 것은 두 개의 컴퍼넌트를 하나로 뭉칠 때 어떤 작업을 해야하는가? 이다.

이는 smaller-to-larger 라는 기법을 알고 있으면 쉽게 떠올릴 수 있다. 즉, 작은뭉치를 큰 뭉치에 넣는다는 것이다. 이 방법에서 임의의 합치는 연산은 전체의 절반 이하의 갯수만 움직이게 된다. 즉, 각 원소는 최대 logM번만 움직이게 된다는 것이다. (이 부분을 처음 접해봤거나 좀 더 자세히 알 고 싶다면 이 글의 윗부분을 읽어보면 좋을 것 같다.) 따라서 합치는 연산을 할 때  두 컴퍼넌트의 간선 저장소의 크기를 확인한다음 작은 쪽의 간선들을 모두 큰 쪽으로 넣어주고 합쳐주면 된다. 총 시간복잡도는 $O(Mlg^2M)$ 인 것 같다. 확실하지는 않지만 시간안에 작동한다는 것은 확신할 수 있다.

코드

#include<bits/stdc++.h>
using namespace std;
typedef long long int ll;
typedef pair<ll,ll> pii;
 
#define ff first
#define ss second
#define ep emplace_back
#define eb emplace
#define pb push_back
 
ll par[1010101];
ll siz[1010101];
ll val[1010101];
 
struct poi{
    ll op,all,half,idx;
};
 
struct com{
    bool operator()(poi &a,poi &b){
        return a.half>b.half;
    }
};
 
struct comp{
    bool operator()(poi &a,poi &b){
        return a.idx>b.idx;
    }
};
 
queue<ll> q;
 
priority_queue<poi,vector<poi>,com> pq[303030];
priority_queue<poi,vector<poi>,comp> fin;
 
ll fi(ll x){
    if(par[x]==x) return x;
    return par[x]=fi(par[x]);
}
 
ll mergee(ll a,ll b){
    a=fi(a);
    b=fi(b);
    if(pq[a].size()<=pq[b].size()){
        par[a]=b;
        val[b]+=val[a];
        while(!pq[a].empty()){
            auto x=pq[a].top();
            pq[a].pop();
            if(fi(x.op)==b) continue;
            pq[b].push(x);
        }
        return b;
    }
    else{
        par[b]=a;
        val[a]+=val[b];
        while(!pq[b].empty()){
            auto x=pq[b].top();
            pq[b].pop();
            if(fi(x.op)==a) continue;
            pq[a].push(x);
        }
        return a;
    }
}
 
 
int main(){
    cin.tie(0);
    ios_base::sync_with_stdio(0);
    ll i,j,k,l,m,n;
    cin>>n>>m;
    if(n==1) while(1);
    for(i=1;i<=n;i++){
        cin>>val[i];
        par[i]=i;
    }
    for(i=1;i<=m;i++){
        ll a,b,c;
        cin>>a>>b>>c;
        if(val[a]+val[b]>=c){
            fin.push({a,b,c,i});
            continue;
        }
        pq[a].push({b,c,val[a]+(c-val[a]-val[b]+1)/2,i});
        pq[b].push({a,c,val[b]+(c-val[a]-val[b]+1)/2,i});
    }
 
    while(!fin.empty()){
        auto x=fin.top(); fin.pop();
        ll a=x.op;
        ll b=x.all;
        if(fi(a)==fi(b)) continue;
        q.push(x.idx);
 
        ll c=mergee(a,b);
 
        while(!pq[c].empty()){
            auto k=pq[c].top();
            k.op=fi(k.op);
            if(fi(c)==fi(k.op)){
                pq[c].pop();
                continue;
            }
            if(val[c]<k.half) break;
            pq[c].pop();
 
            if(k.all<=val[c]+val[k.op]){
                fin.push({c,k.op,k.idx,k.idx});
            }
            else{
                pq[k.op].push({c,k.all,val[k.op]+(k.all-val[k.op]-val[c]+1)/2,k.idx});
                pq[c].push({k.op,k.all,val[c]+(k.all-val[k.op]-val[c]+1)/2,k.idx});
            }
        }
    }
 
    cout<<q.size()<<'\n';
    while(q.size()){
        cout<<q.front()<<' ';
        q.pop();
    }
 
 
}