[c++] 백준 #3830 교수님은 기다리지 않는다.

< 풀이 >

맨처음엔 정말 naive하게 O(N^2) dfs 풀이로 진행했다가 TLE를 맛보고 union disjoint set을 이용해서 풀어야 한다는 것을 알게 되었다.

b가 a보다 무겁다는 걸 그래프로 나타내면 아래와 같다. b가 무거우니까 diff[b] - k = diff[a]가 된다.

 

b가 a보다 k그램 무겁다는 걸 그래프로 나타낸 것

 

여기서 b와 a에게 부모가 있다는 걸 생각하고 그래프를 그려보자.

b의 부모에 a의 부모를 붙일거면 diff[b] - k = x + diff[a]가 성립한다. 따라서 x = diff[b] - diff[a] - k가 된다.

반대로 a의 부모에 b의 부모를 붙일 거면 -x를 붙여주면 된다. 

a, b가 부모가 있는 경우 그래프 갱신 방법

 

** 각 테케마다 p, diff, n을 초기화해줘야하는데 여기서 실수로 for문 2개를 넣어버려서 이상한 곳에서 TLE가 나서 나의 소중한 정답률 하락에 기여했다. ㅠㅠ

 

< 소스코드 >

#include <bits/stdc++.h>
using namespace std;
#define INF 987654321
typedef long long ll;
typedef pair<int, int> ii;
typedef tuple<int, int, int> iii;

#define pb push_back
#define N_MAX 100010

ll p[N_MAX]; // 누가 부모인지.
ll diff[N_MAX]; // 부모와의 거리
ll n[N_MAX]; // set의 size

ll my_find(ll a){
    if (p[a] == a) return a;
    else {
        ll prev = my_find(p[a]); // 제일 위에 루트 찾기
        diff[a] += diff[p[a]]; // 나는 내 위에 애 w에 내 w를 더해주고, root와 연결해줌.
        return p[a] = prev;
    }
}

// b가 a보다 k그램 더 무겁다.
// a를 b에 붙이는 게 default.
void my_union(ll a, ll b, ll k){
    ll a_root = my_find(a);
    ll b_root = my_find(b);
    if (a_root != b_root){
        if (n[b_root] >= n[a_root]){
            // a를 b에 붙인다.
            p[a_root] = b_root;
            ll new_diff = diff[b] - diff[a] - k;
            diff[a_root] += new_diff;
            n[b_root] += n[a_root];
            n[a_root] = 1; 
        }
        else { // b를 a에 붙이고 diff는 음수로 붙여준다.
            p[b_root] = a_root;
            ll new_diff = diff[b] - diff[a] - k;
            diff[b_root] -= new_diff;
            n[a_root] += n[b_root];
            n[b_root] = 1; 
        }
    }
}

int main(){
    ios::sync_with_stdio(false);
    cin.tie(0);
    while(1){
        ll N, M; cin >> N >> M;
        for(int i=0; i<=N;i++){
            p[i] = i;
            diff[i] = 0;
            n[i] = 1;
        }
        if (N==0 && M==0) break;
        for(ll i=0;i<M;i++){
            char q; cin >> q;
            if (q== '!'){
                ll a,b,w; cin >> a >> b >> w; // b가 a보다 w만큼 무겁다.
                my_union(a, b, w);
            }
            else {
                ll a, b; cin >> a >> b;
                ll a_root = my_find(a);
                ll b_root = my_find(b);
                if (a_root == b_root){
                    cout << diff[b] - diff[a] << '\n';
                }
                else {
                    cout << "UNKNOWN\n";
                }
            }
        }
    }
    return 0;
}

 

 

참고 블로그

https://lyzqm.blogspot.com/2018/06/3830.html