물대기 [BOJ1368]

문제

선주는 자신이 운영하는 N개의 논에 물을 대려고 한다. 물을 대는 방법은 두 가지가 있는데 하나는 직접 논에 우물을 파는 것이고 다른 하나는 이미 물을 대고 있는 다른 논으로부터 물을 끌어오는 법이다. 각각의 논에 대해 우물을 파는 비용과 논들 사이에 물을 끌어오는 비용들이 주어졌을 때 최소의 비용으로 모든 논에 물을 대는 것이 문제이다.

풀이

얼핏보면 각 우물을 직접파기 or 다른 우물에서 길어오기 와 같은 이분법적사고로 착각하기 쉽다. 하지만 우물을 직접파는걸 원천우물(이라는 가상의 우물을 만들고)에서 길어오는 것으로 편입시키면 하나의 방법으로 통일된다.

이제 문제는 모든 논(원천가상우물포함)을 최소비용으로 연결하면 된다. 이건 결국 최소스패닝트리문제이다.

code

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;

// 간선을 표현하는 구조체
struct Edge {
    int src, dest, weight;
};

// 유니온-파인드 알고리즘을 위한 함수들
int find(vector<int>& parent, int i) {
    if (parent[i] == i) return i;
    return find(parent, parent[i]);
}

void Union(vector<int>& parent, int x, int y) {
    int xset = find(parent, x);
    int yset = find(parent, y);
    if(xset < yset) parent[xset] = yset;
    if(xset > yset) parent[yset] = xset;
}

// 크루스칼 알고리즘
int kruskalMST(vector<Edge>& edges, int n) {
    int mst_weight = 0;

    // 간선들을 가중치 기준으로 정렬
    sort(edges.begin(), edges.end(), [](const Edge &a, const Edge &b) {
        return a.weight < b.weight;
    });

    vector<int> parent(n + 1);
    for (int i = 1; i <= n; ++i) parent[i] = i;

    for (Edge &edge : edges) {
        int x = find(parent, edge.src);
        int y = find(parent, edge.dest);

        if (x != y) {
            mst_weight += edge.weight;
            Union(parent, x, y);
        }
    }

    return mst_weight;
}

int main() {
	ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    int n,weight;
    cin >> n;

    vector<Edge> edges; // 0:원천우물 1~n: 논
    for (int i = 0; i < n; ++i) {
    	Edge newEdge; 
    	cin >> newEdge.weight;
    	newEdge.src = i+1; newEdge.dest = 0;
    	edges.push_back(newEdge);
    }
    for (int i = 1 ; i <= n; ++i) {
    	for (int j = 1; j <= n ; ++j){
	    	Edge newEdge; 
	    	cin >> newEdge.weight;
	    	newEdge.src = i; newEdge.dest = j;
	    	edges.push_back(newEdge);
    	}
    }	
    

    cout << kruskalMST(edges, n) << endl;

    return 0;
}