백준 15481번: 그래프와 MST (C++)

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

15481번: 그래프와 MST

---

Solved By: MST, LCA

 

MST를 구합니다. 그리고, 주어진 Query Node 사이에 maxDist를 빼고, Query Edge의 w를 더하는 것이 이번 문제의 핵심입니다.

여기서 궁금한 점은 경로를 맘대로 빼고, 더하는 행위가 component가 나뉘지 않는지 걱정했지만 몇 개의 case를 그려보면 그렇지 않다는 것을 알 수 있습니다.

 

maxDist를 구하는 방법은 LCA(using Sparse Table)를 사용합니다

#include <iostream>
#include <vector>
#include <algorithm>
#define MAX 200001
#define LOG_MAX 18
#define ll long long 
#define pii pair<int, int>
using namespace std;

vector<pair<int, ll>> adj[MAX]; // MST-TREE
vector<pair<pii, ll>> edges;
vector<pair<pii, ll>> edges2; // Sort 때문에 기존 edges를 저장해놓는 temp가 필요
int parent[MAX][LOG_MAX];
ll dist[MAX][LOG_MAX];
int level[MAX];
int uf[MAX];
int n, m;
ll MST_W = 0;

int getRoot(int node){
    if(node == uf[node]) return node;
    return uf[node] = getRoot(uf[node]);
}

bool find(int a, int b){
    int ga = getRoot(a);
    int gb = getRoot(b);
    
    if(ga != gb) return true;
    else return false;
}

void unionNodes(int a, int b){
    int ga = getRoot(a);
    int gb = getRoot(b);
    
    if(ga < gb) uf[gb] = ga;
    else uf[ga] = gb;
}

bool cmp(pair<pii, ll> a, pair<pii, ll> b){
    return a.second < b.second;
}

void dfs(int now, int par){
    for(int i = 0; i < adj[now].size(); i++){
        int next = adj[now][i].first;
        ll nextCost = adj[now][i].second;
        
        if(next == par) continue;
        
        level[next] = level[now] + 1;
        parent[next][0] = now;
        dist[next][0] = nextCost;
        dfs(next, now);
    }
}

void swap(int* a, int* b){
    int* temp = a;
    a = b;
    b = temp;
}

ll maxDist(int a, int b){
    if(level[a] < level[b]) swap(a, b);
    
    ll ret = 0;
    int diff = level[a] - level[b];
    
    for(int i = LOG_MAX - 1; i >= 0; i--){
        if(diff >= 1 << i){
            diff -= 1 << i;
            ret = max(ret, dist[a][i]);
            a = parent[a][i];
        }
    }
    
    if(a != b){
        for(int i = LOG_MAX - 1; i >= 0; i--){
            if(parent[a][i] != 0 && parent[a][i] != parent[b][i]){
                ret = max(ret, dist[a][i]);
                ret = max(ret, dist[b][i]);
                
                a = parent[a][i];
                b = parent[b][i];
            }
        }
        ret = max(ret, dist[a][0]);
        ret = max(ret, dist[b][0]);
    }
    
    return ret;
}

int main(){
    cin >> n >> m;
    for(int i = 0; i < m; i++){
        int a, b; ll w;
        cin >> a >> b >> w;
        edges.push_back({{a, b}, w});
        edges2.push_back({{a, b}, w});
    }
    
    // Union-Find init
    for(int i = 1; i <= n; i++) uf[i] = i;
    sort(edges.begin(), edges.end(), cmp);
    
    for(int i = 0; i < edges.size(); i++){
        int first = edges[i].first.first;
        int second = edges[i].first.second;
        if(find(first, second)){
            unionNodes(first, second);
            adj[first].push_back({second, edges[i].second});
            adj[second].push_back({first, edges[i].second});
            MST_W += edges[i].second;
        }
    }
    
    // LCA init
    dfs(1, -1);
    
    // MST-TREE LCA's parent init
    for(int j = 1; j < LOG_MAX; j++){
        for(int i = 1; i <= n; i++){
            parent[i][j] = parent[parent[i][j - 1]][j - 1];
            dist[i][j] = max(dist[i][j - 1], dist[parent[i][j - 1]][j - 1]);
        }
    }
    
    // query
    // 
    for(int i = 0; i < edges2.size(); i++){
        int first = edges2[i].first.first;
        int second = edges2[i].first.second;
        //cout << "maxDist(first, second): " << maxDist(first, second) << '\n';
        cout << MST_W - maxDist(first, second) + edges2[i].second << '\n';
    }
}