1168. 水資源分配優化

思路

會員題

村裏面一共有 n 棟房子。我們希望通過建造水井和鋪設管道來爲所有房子供水。
對於每個房子 i，我們有兩種可選的供水方案：
一種是直接在房子內建造水井，成本爲 wells[i]；
另一種是從另一口井鋪設管道引水，數組 pipes 給出了在房子間鋪設管道的成本，其中每個 pipes[i] = [house1, house2, cost] 代表用管道將 house1 和 house2 連接在一起的成本。當然，連接是雙向的。
請你幫忙計算爲所有房子都供水的最低總成本。

鋪設管道時，對於一間房子來說，可以從鄰居拉水管過來，或是在自己這個點建造。
看似要對最小生成樹修改很多內容，其實不用。
在房子當中建立井的過程，我們可以想像成：從一個代表水源的虛擬節點，連到各自房子當中。
因此，題目就轉變成了，所有房子，以及虛擬水源互相連接，所需要的最小成本，套用最小生成樹即可求解。

由於是會員題，以下程式碼沒有實際提交過

程式碼

最小生成樹 Kruskal

#include <iostream>
#include <vector>
#include <algorithm>
#include <numeric>
using namespace std;
struct Edge {
    int u, v, w;
    Edge(int u, int v, int w) : u(u), v(v), w(w) {}
    Edge(const vector<int>& vec) : u(vec[0]), v(vec[1]), w(vec[2]) {}
    bool operator<(const Edge& other) {
        return this->w < other.w;
    }
};
class Solution {
private:
    static const int MX = 1e4+5;
    int parent[MX];
    int find(int x) {
        if(parent[x] != x) {
            parent[x] = find(parent[x]);
        }
        return parent[x];
    }
    bool isSameSet(int x, int y) {
        return find(x) == find(y);
    }
    void unite(int x, int y) {
        parent[find(x)] = parent[find(y)];
    }
public:
    int minCostToSupplyWater(int n, vector<int>& wells, vector<vector<int>>& pipes) {
        // 節點 0 代表水源，一個虛擬節點。
        iota(parent, parent + n + 1, 0);
        vector<Edge> edges;
        for(int i = 1; i <= n; i++) {
            edges.push_back(Edge(0, i, wells[i - 1]));
        }
        for(auto& vec : pipes) {
            edges.push_back(Edge(vec));
        }
        sort(edges.begin(), edges.end());
        int res = 0;
        int edgeCnt = 0;
        for(Edge& edge : edges) {
            if(!isSameSet(edge.u, edge.v)) {
                res += edge.w;
                unite(edge.u, edge.v);
                if(++edgeCnt == n) { // 共 n + 1 個點
                    break;
                }
            }
        }
        return res;
    }
};
int main(void) {
    Solution sol;
    int n = 3;
    vector<int> wells = {1, 2, 2};
    vector<vector<int>> pipes = {{1, 2, 1}, {2, 3, 1}};
    cout << sol.minCostToSupplyWater(n, wells, pipes);
}

複雜度分析

• 時間複雜度：$O(m\log{m})$
• 空間複雜度：$O(n+m)$