프로그래밍/코딩테스트

[프로그래머스/크루스칼 알고리즘] 섬 연결하기

바토파 2024. 9. 19. 21:49
반응형

문제 링크

https://school.programmers.co.kr/learn/courses/30/lessons/42861


문제 설명


코드 전문

#include <string>
#include <vector>
#include <algorithm>

using namespace std;

int get_parent(int i, vector<int>& parent){
    if (parent[i] == i) return i;
    return parent[i] = get_parent(parent[i], parent);
}

void union_parent(int a, int b, vector<int>& parent){
    int a_parent = get_parent(a, parent);
    int b_parent = get_parent(b, parent);
    
    if (a_parent < b_parent) parent[b_parent] = a_parent;
    else parent[a_parent] = b_parent;
}

bool compare(vector<int> a, vector<int> b){
    return a[2] < b[2];
}

int solution(int n, vector<vector<int>> costs) {
    int answer = 0;

    sort(costs.begin(), costs.end(), compare);
    
    vector<int> parents(n);
    
    for (int i = 0; i < n; i++){
        parents[i] = i;
    }
    
    for (auto c : costs){
        if (get_parent(c[0], parents) != get_parent(c[1], parents)){
            union_parent(c[0], c[1], parents);
            answer += c[2];
        }
    }
    
    return answer;
}

풀이

크루스칼 알고리즘을 사용해서 풀면 된다.

 

i의 부모를 찾는 함수이다.

parent[i] = i 면 i가 root라는 뜻이기 때문에 i를 반환한다.

아니라면 parent[i]의 부모 노드를 재귀 호출을 통해 찾는다.

int get_parent(int i, vector<int>& parent){
    if (parent[i] == i) return i;
    return parent[i] = get_parent(parent[i], parent);
}

 

a와 b의 부모를 합치는 함수이다.

a의 부모의 b의 부모를 찾고 부모의 크기가 작은 것으로 합친다.

 

void union_parent(int a, int b, vector<int>& parent){
    int a_parent = get_parent(a, parent);
    int b_parent = get_parent(b, parent);
    
    if (a_parent < b_parent) parent[b_parent] = a_parent;
    else parent[a_parent] = b_parent;
}

 

정렬을 위한 함수이다.

cost를 기준으로 오름차순 정렬한다.

 

bool compare(vector<int> a, vector<int> b){
    return a[2] < b[2];
}

 

부모를 기록할 배열을 하나 생성해 처음에는 자기 자신을 부모로 정해준다.

그리고 costs를 순회하며 root가 하나만 존재할 수 있도록 합쳐준다.

마지막으로 answer에 cost를 더한다.

int solution(int n, vector<vector<int>> costs) {
    int answer = 0;

    sort(costs.begin(), costs.end(), compare);
    
    vector<int> parents(n);
    
    for (int i = 0; i < n; i++){
        parents[i] = i;
    }
    
    for (auto c : costs){
        if (get_parent(c[0], parents) != get_parent(c[1], parents)){
            union_parent(c[0], c[1], parents);
            answer += c[2];
        }
    }
    
    return answer;
}
반응형