백준 16562번: 친구비 C++코드(Union Find, 유니온 파인드, 분리 집합)

#include<iostream>

using namespace std;

int friendCost[10001];
int friendCostMin[10001] = {0,};
int parents[10001];

int find(int n){
    if(parents[n-1] == n){
        return n;
    }else{
        int temp = find(parents[n-1]);
        return temp;
    }
}

void Union(int a, int b){
    a = find(a);
    b = find(b);

    if(a < b){
        parents[b-1] = a;
    }else{
        parents[a-1] = b;
    }
}

void Union2(int n){
    int now = n;
    n = find(n);
    parents[now-1] = n;
}

int main(){
    int n,m,k,f1,f2;
    cin >> n >> m >> k;
    for(int i=0; i<n; i++){
        parents[i] = i+1;
    }
    for(int i=0; i<n; i++){
        cin >> friendCost[i];
    }

    for(int i=0; i<m; i++){
        cin >> f1 >> f2;
        Union(f1, f2);
    }


    for(int i=1; i<=n; i++){
        Union2(i);
    }

    for(int i=0; i<n; i++){
        if(friendCostMin[parents[i]-1] == 0){
            friendCostMin[parents[i]-1] = friendCost[parents[i]-1];
        }else if(friendCostMin[parents[i]-1] > friendCost[i]){
            friendCostMin[parents[i]-1] = friendCost[i];
        }
    }
    
    int start = k;

    for(int i=0; i<n; i++){
        if(friendCostMin[i] != 0){
            k -= friendCostMin[i];
        }
    }
    if(k < 0){
        cout << "Oh no";
    }else{
        cout << start-k;
    }
   
}

 

오랜만에 유니온 파인드를 풀어봤다. 그래서 엄청 어려운 로직이 필요하지 않음에도 좀 시간이 걸렸다. 코드가 아직도 좀 난잡하기도 하다. Union함수를 Find함수에 그냥 포함시킨 상태에서 해결하다가 돈 계산하는 부분을 만들고, Union 함수를 다시 만들어서 Union2 함수를 따로 만들었다.

 

로직 설명:

 

입력을 받으면서 한번 Union 함수를 통해, 기본적인 연결을 해준 후, Union2함수를 통해 전부 같은 부모로 묶어주는 식으로 구현했다. 이후 for문에서 friendCostMin 배열의 [parents[i]-1]에 아무것도 할당된 적이 없다면 friendCostMin 배열의 [parents[i]-1]를 할당해준다. 그리고 할당된 적이 있다면 friendCostMin[parents[i]-1]와 friendCost[parents[i]-1]를 비교하여 더 작은 값을 할당한다. 이 과정을 거친 friendsCostMin 배열에서 0이 아니라면 할당된 적이 있는 인덱스이기 때문에 그 값을 K(돈)에서 빼준다. 이후 K가 0보다 작다면 돈이 모자란 것이고 0보다 크거나 같다면 (처음 가진 돈-친구비를 낸 돈 = 최소비용)최소비용을 출력한다.