[c++] 백준 16950 레드 블루 스패닝 트리 2

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

16950번: 레드 블루 스패닝 트리 2

이 문제를 풀기 위해서는 (백준 4792 레드 블루 스패닝 트리)의 내용을 꼭 알고 있어야 한다.

위 링크의 글에서 설명했듯이 파란 간선이 t개인 mst에서 t+1인 mst를 만들려면

정점 u, v를 잇는 사용되지 않은 파란색 간선을 찾고,

스패닝 트리에서 u에서 v로 가는 경로중 빨간색 간선을 지우고 파란색 간선을 추가하면 된다.

이를 그대로 구현하면 된다.

 

경로를 구해주는 것이 생각보다 까다롭다.

정점의 수가 최대 1,000개이기 때문에 구현의 방법은 많을 것 같다.

LCA를 구하는 과정에서 빨간색 간선을 찾을 수 있다.(희소 배열을 사용하지 않고 O(N)에 해결 가능)

그냥 DFS를 이용하여 경로를 찾는 방법도 있다.

 

그저 구현만 되면 풀리는 문제이기 때문에 풀이법은 자유롭다.

 

나의 코드 구조는 이렇다.

1. 크루스칼로 파란색 간선이 최소로 사용된 mst를 만든다.

2. 사용되지 않은 파란색 간선과 사용된 빨간색 간선의 정보를 저장.

3. 사용되지 않은 파란색 간선을 고르고, 연결할 정점들을 dfs로 경로를 구한다

4. 경로에 빨간색 간선이 있는가를 체크 후 있으면 빨간색 간선을 삭제, 파란색 간선 추가

5. 이미 조건을 만족한 mst면 출력 아니면 사용하지 않은 파란색 간선 사용.

6. 사용하지 않은 모든 파란색 간선을 사용하면서 조건을 만족한 mst면 출력

7. 사용하지 않은 모든 파란색 간선을 전부 사용했음에도 조건을 만족하지 못하면 0 출력

 

※주의사항

생각보다 시간이 빡빡하다.

mst를 인접 리스트로 구현했을 때는 약 100ms를 받았지만, 인접 행렬로 구현했을 때는 약 700ms를 받았다.

문제 시간제한이 1000ms인 것을 생각해 보면 아슬아슬하다.

인접 행렬을 사용하는 것이 변수도 덜 사용하고 구현 난이도도 쉬운 것 같아서 인접 행렬로 코드를 작성하기는 했다.

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
int n,m,k,used_blue;
vector<tuple<char,int,int>> edge;
vector<pair<int,int>> unused_blue;
int road[1010],parent[1010];
int used_red[1010][1010];
int mst[1010][1010];
int FIND(int x){
    if(parent[x] == x) return x;
    return parent[x] = FIND(parent[x]);
}
pair<int,int> dfs(int cur, int target, int par, int idx){
    road[idx] = cur;
    if(cur == target){
        // 경로 중에 빨간색 간선 찾기
        for(int i=1; i<=idx; i++)
            // 있으면
            if(used_red[road[i]][road[i-1]])
                return{road[i],road[i-1]};
        // 없으면
        return {0,0};
    }
    // dfs 돌기
    int p,q;
    for(int next=1; next<=n; next++){
        if(mst[cur][next]){
            if(next == par) continue;
            tie(p,q) = dfs(next,target,cur,idx+1);
            if(p && q) return {p,q};
        }
    }
    return {0,0};
}
void OUTPUT(){
    for(int i=1; i<=n; i++)
        for(int j=i+1; j<=n; j++)
            if(mst[i][j])
                cout << i << ' ' << j << '\n';
}
int main(){
    ios::sync_with_stdio(0);
    cin.tie(0);
    // 입력
    cin>>n>>m>>k;
    for(int i=0; i<=n; i++)
        parent[i] = i;
    for(int i=0; i<m; i++){
        char c; int u,v;
        cin>>c>>u>>v;
        edge.push_back({c,u,v});
    }
    // 크루스칼
    sort(edge.rbegin(),edge.rend());
    for(auto [c,u,v] : edge){
        int fu=FIND(u), fv=FIND(v);
        if(fu != fv){
            // 인접 행렬로 mst 구현
            mst[u][v]=mst[v][u]=1;
            parent[fu] = fv;
            // 사용할 정보들 기록해 놓기
            // used_red, used_blue, unused_blue
            if(c == 'R') used_red[u][v]=used_red[v][u]=1;
            else used_blue++;
        }
        else{
            if(c == 'B')
                unused_blue.push_back({u,v});
        }
    }
    // mst의 파란색 간선을 증가시키기
    if(used_blue == k){OUTPUT();return 0;}
    for(auto [u,v] : unused_blue){
        int ru,rv;
        tie(ru,rv) = dfs(u,v,0,0);
        // 경로에 빨간색 간선이 없으면 continue
        if(!ru && !rv) continue;
        // 있으면 mst 수정
        mst[ru][rv]=mst[rv][ru]=0;
        mst[u][v]=mst[v][u]=1;
        used_red[ru][rv]=used_red[rv][ru]=0;
        used_blue++;
        // 조건에 맞는 mst이면 출력
        if(used_blue == k){OUTPUT();return 0;}
    }
    cout << 0;
    return 0;
}