[BOJ 3830] 교수님은 기다리지 않는다 (C++)

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

3830번: 교수님은 기다리지 않는다

#include <bits/stdc++.h>
#define ll long long
using namespace std;

int parents[100001];
ll diff[100001];    //root 와의 차이

int find(int a){
    if (parents[a]==a){
        return a;
    }
    int p=find(parents[a]);
    diff[a]+=diff[parents[a]];      //갱신
    return parents[a]=p;
}

void union_(int a, int b, ll w){
    int pa=find(a);
    int pb=find(b);
    if (pa!=pb){
        diff[pb]=diff[a]-diff[b]+w;
        parents[pb]=pa;
    }
}

int main(){
    int N, M;
    char c;
    int a,b;
    ll w;
    vector<string> ans;
    cin>>N>>M;
    while (N!=0 || M!=0){
        for (int i=1; i<=N; i++){
            parents[i]=i;
            diff[i]=0;
        }
        for (int i=0; i<M; i++){
            cin>>c;
            if (c=='!'){
                cin>>a>>b>>w;
                union_(a,b,w);
            }
            else if (c=='?'){
                cin>>a>>b;
                if (find(a)!=find(b)){
                    ans.push_back("UNKNOWN");
                }
                else{
                    ans.push_back(to_string(diff[b]-diff[a]));
                }
            }
        }
        cin>>N>>M;
    }
    for (int i=0; i<ans.size(); i++){
        cout<<ans[i]<<"\n";
    }
    return 0;
}

2주 전에 특강에서 풀었던 문젠데 그때도 틀리고 지금 고민해봐도 안풀려서 해답 보고 풀었다. 플래티넘 넘사..

 

union-find 문제.. 근데 재귀를 곁들인..

 

diff는 root 노드와의 차이를 나타낸다. 

A 노드의 부모가 B 노드고, B 노드가 루트 노드인(자기 자신이 부모인) C를 가리키고 있을 때, 

A의 root 노드는 C가 된다. 근데 현재 diff[A]는 B와의 거리를 나타내므로, 부모 노드가 다른 노드를 부모로 삼고 가리키고 있을 때 같은 그룹에 속해서 차이를 알 수 있도록 해야된다. diff[B]를 diff[A]에다 더해서 루트 노드와의 최종 거리를 구해줘야 되는게 포인트다. find 함수에서 이상의 과정을 거치면서 루트 노드와의 차이가 diff 에 저장되며 갱신된다.

 

union 함수는 a노드는 pa와 diff[a] 만큼 차이, b 노드는 pb와 diff[b] 만큼 차이난다고 했을 때 pa와 pb 간의 거리는 다음과 같이 구해준다. 

pb의 루트는 이제 pa가 되었으므로, pb와 B는 -diff[B] 만큼 차이, B와 A는 w만큼 차이, A와 pa는 diff[A] 만큼 차이나므로 이를 모두 더해주면 diff[A]-diff[B]+w가 된다. 이 아이디어를 떠올리는게 힘든것같다..

 

pb 기준으로 diff 써진 애들은 나중에 find 호출하게 되면 pa의 루트인 pa 기준으로 다시 갱신된다.