백준 17398번. 통신망 분할(c++) / union-find

이 문제는 처음에 딱 읽었을 때 bfs로 풀면 간단하겠다라고 생각했지만, 통신탑과 간선의 개수가 각각 최대 100,000 인것을 보고

bfs로 풀다가는 백퍼 시간초과가 나겠구나 싶었다.

 

그리고 다른 방법을 고민하다가 도저히 아이디어가 생각나지 않아, 문제 유형을 검색해보니

바로 "유니온파인드"를 이용하는 문제였다.

 

끊을 간선을 제외한 모든 간선을 먼저 연결해주고 끊을 간선들을 역순으로 하나씩 연결해주며

값을 계산해나가면 된다.

예를 들어, a->b a를 b와 연결해준다면 a의 부모 노드를 찾고 b의 부모노드를 찾아서 연결해주고

b의 부모노드의 값에 a의 부모노드의 값을 더해주면 된다. (cost[b] += cost[a] 이런식으로 구현했다.)

 

코드 첨부

// 통신망 분할 (union-find)
#include <iostream>
#include <vector>
#include <map>
using namespace std;
 
int n, m, q, x, y, z;
 
vector<int> disconnect; // 끊을 곳
vector<pair<int, int>> v;
map<int, int> M;
int connect[100001] = {0, };
long long cost[100001] = {0, };
long long total = 0;
 
int find(int a){
    if(a==connect[a]){
        return a;
    }else{
        return connect[a] = find(connect[a]);
    }
}
 
 
int main()
{
    cin>>n>>m>>q;
    for(int i=1; i<=m; i++){
        cin>>x>>y;
        v.push_back(make_pair(x, y));
    }
    for(int i=0; i<q; i++){
        cin>>z;
        M[z] = 1;
        disconnect.push_back(z); // z번째 연결 요소를 끊는다.
    }
 
    for(int i=1; i<=n; i++){
        cost[i] = 1;
        connect[i] = i;
    }
 
    for(int i=0; i<v.size(); i++){
        if(M[i+1]==1){ // 끊어야 할 연결요소
            continue;
        }
        int a = find(v[i].first);
        int b = find(v[i].second);
 
        if(a!=b){
            // a->b
            connect[a] = b;
            cost[b] += cost[a];
        }
    }
 
    // 거꾸로 연결해준다.
    for(int i=disconnect.size()-1; i>=0; i--){
        int idx = disconnect[i]-1;
        int a = find(v[idx].first);
        int b = find(v[idx].second);
 
        if(a!=b){
            connect[a] = b;
            total += cost[a]*cost[b];
            cost[b] += cost[a];
        }
    }
 
    cout<<total<<"\n";
    return 0;
}

문제 첨부

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

17398번: 통신망 분할