백준 1368 물대기 (C++)

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

#include <iostream>
#include <vector>
#include <algorithm>

using namespace std;

vector<int> sCost;
vector<vector<int>> cCost;

int parent[301];

int findP(int k) {
    if (parent[k] == k) {
        return k;
    } else {
        return parent[k] = findP(parent[k]);
    }
}

void uni(int a, int b){

    int pA = findP(a);
    int pB = findP(b);

    if (pA == pB) {
        return;
    }

    if (pA > pB) {
        parent[pA] = pB;
    } else {
        parent[pB] = pA;
    }
}

int main() {

    ios::sync_with_stdio(false);
    cin.tie(NULL);
    cout.tie(NULL);

    int n;
    cin >> n;
    int ans = 0;

    sCost.resize(n);

    for (int i = 0; i < n; i++) {
        cin >> sCost[i];
        parent[i] = i;
    }
    parent[n] = n;

    cCost.resize(n, vector<int>(n));

    for (int i = 0; i < n; i++) {
        for (int j = 0; j < n; j++) {
            cin >> cCost[i][j];
        }
    }

    vector<pair<int, pair<int, int>>> cost;

    for (int i = 0; i < n; i++) {
        cost.push_back(make_pair(sCost[i], make_pair(i, n)));
    }

    for (int i = 0; i < n - 1; i++) {
        for (int j = i + 1; j < n; j++) {
            cost.push_back(make_pair(cCost[i][j], make_pair(i, j)));
        }
    }


    sort(cost.begin(), cost.end());

    int cnt = 0;
    for (int i = 0; i < cost.size(); i++) {
        if (cnt == n) {
            break;
        }

        if (findP(cost[i].second.first) != findP(cost[i].second.second)) {
            uni(cost[i].second.first, cost[i].second.second);
            ans += cost[i].first;
            cnt++;
        }
    }



    cout << ans;

    return 0;
}

1. 문제

n개의 논에 물을 대려 한다. 물을 채우는 방법은 직접 논에 우물을 파는 방법과 논 끼리 연결해서 다른 논에 있는 물을 끌어오는 방법이다. 직접 논에 우물을 파는 데 드는 비용과 각 논끼리 연결하는데 드는 비용이 주어질 때, 모든 논에 물을 대는 데 드는 최소 비용을 구하라.

2. 풀이

MST 문제인데, 아이디어가 중요한 문제였다. 이 문제는 직접 우물을 파는 것과 다른 논의 물을 끌어오는 것을 비교하는 문제이다. 중요한 것은 모든 논에 물을 대려면 적어도 하나의 논은 직접 우물을 파는 방식을 선택해야 한다는 것이다. 

이를 하나의 노드로 생각하는 방법이다. 직접 우물을 파는 경우를 하나의 가상 노드로 취급하고, 각 노드와 연결하는 데, 우물을 파는 비용을 입력한다.

4
5
4
4
3
0 2 2 2
2 0 3 3
2 3 0 4
2 3 4 0

이런 입력이 주어졌다면

이런 식으로 연결을 하는 것이다. 반드시 한 논은 우물을 파야 하고, 최소 비용으로 연결을 해야 하니 이부터는 일반적인 MST 알고리즘을 사용하면 풀 수 있다.

int parent[301];

int findP(int k) {
    if (parent[k] == k) {
        return k;
    } else {
        return parent[k] = findP(parent[k]);
    }
}

void uni(int a, int b){

    int pA = findP(a);
    int pB = findP(b);

    if (pA == pB) {
        return;
    }

    if (pA > pB) {
        parent[pA] = pB;
    } else {
        parent[pB] = pA;
    }
}

일단 kruskal 알고리즘에 필요한 부모 노드 찾는 함수, union 연산 함수를 구현해준다.

    int n;
    cin >> n;
    int ans = 0;

    sCost.resize(n);

    for (int i = 0; i < n; i++) {
        cin >> sCost[i];
        parent[i] = i;
    }
    parent[n] = n;

    cCost.resize(n, vector<int>(n));

    for (int i = 0; i < n; i++) {
        for (int j = 0; j < n; j++) {
            cin >> cCost[i][j];
        }
    }

각각 우물을 파는 데 드는 비용, 두 논을 연결하는 데 드는 비용을 구해준다.

 vector<pair<int, pair<int, int>>> cost;

    for (int i = 0; i < n; i++) {
        cost.push_back(make_pair(sCost[i], make_pair(i, n)));
    }

    for (int i = 0; i < n - 1; i++) {
        for (int j = i + 1; j < n; j++) {
            cost.push_back(make_pair(cCost[i][j], make_pair(i, j)));
        }
    }

<pair<cost, pair<u, v>>>의 자료구조를 만들고, 가상의 노드 n을 만들어 n과는 우물을 파는 데 드는 비용을, 각 논 끼리는 연결하는 데 드는 비용을 넣어준다.  a-b를 연결하는 비용과 b-a를 연결하는 비용은 같으므로 하나만 넣어준다.

 sort(cost.begin(), cost.end());

    int cnt = 0;
    for (int i = 0; i < cost.size(); i++) {
        if (cnt == n) {
            break;
        }

        if (findP(cost[i].second.first) != findP(cost[i].second.second)) {
            uni(cost[i].second.first, cost[i].second.second);
            ans += cost[i].first;
            cnt++;
        }
    }



    cout << ans;

이후 kruskal 알고리즘을 구현한다. 가상 노드를 포함하여 n+1개의 노드가 있으므로 총 edge의 수는 n개여야 한다. 부모 노드를 비교하여 cycle이 만들어지지 않으면 해당 edge를 붙이고, 그 비용을 추가해준다. 

 연산 후 총 비용 ans를 출력하면 정답을 구할 수 있다.