본문 바로가기
Programming/Algorithm(알고리즘)

최소신장트리 (MST) 핵심 코드 | c++

by 악마근육남 2021. 2. 19.
반응형

최소신장트리의 가장 기본이 되는 코드입니다.

해당 유형의 문제들은 다 이 코드를 수정해서 풀이할 수 있을 것입니다.

사실 최소신장트리의 핵심은 union find 이기 때문에 최소신장트리라고 지칭하는것 보다 union find 응용 문제라고 하는것이 더 나을지 모릅니다. 

int parent[10001];

struct node
{
    int start;
    int end;
    int dist;
    
    node(int a, int b, int c)
    {
        start = a;
        end = b;
        dist = c;
    }
};

struct compare
{
  bool operator()(node a, node b)
    {
        return a.dist > b.dist;
    }
};

priority_queue<node,vector<node>, compare> v;

int find(int u)
{
    if(u == parent[u]) return u;
    
    return parent[u] = find(parent[u]);
}

void mergeD(int u, int v)
{
    u = find(u);
    v = find(v);
    
    parent[v] = u;
}


int main(int argc, const char * argv[]) {
  
    int V, E;
    int ans = 0;
    scanf("%d %d", &V, &E);
    
    for (int i = 0; i < V; i++) {
        parent[i] = i;
    }
    
    for (int i = 0; i < E; i++) {
        int a, b, c;
        
        scanf("%d %d %d", &a, &b, &c);
        
        v.push(node(a,b,c));
    }
    
    while(!v.empty())
    {
        if(find(v.top().start) != find(v.top().end))
        {
            mergeD(v.top().start, v.top().end);
            ans += v.top().dist;
        }
        
        v.pop();
    }
    
    
    cout << ans << endl;
    
    return 0;
}

메인 코드를 설명하자면

들어오는 노드에 대한 정보를 소모성 비용이 적은 순서대로 정렬해서 받아줍니다. 

그럼 그 순서대로 두개의 노드가 붙어 있지 않다면 이어주며 비용을 더해가면 최소비용으로 전체를 이어줄 수 있다는 개념입니다.

그리고 다른 문제들은 전부 이개념에서 변주 입니다.

입력 부분과 union find 함수를 제외하고는 5줄 밖에 안되는 너무도 짧은 코드이기 때문에 숙지하기 쉬울것 같습니다.

반응형