1168. Optimize Water Distribution in a Village

There are n houses in a village. We want to supply water for all the houses by building wells and laying pipes.

For each house i, we can either build a well inside it directly with cost wells[i - 1] (note the -1 due to 0-indexing), or pipe in water from another well to it. The costs to lay pipes between houses are given by the array pipes where each pipes[j] = [house1j, house2j, costj] represents the cost to connect house1j and house2j together using a pipe. Connections are bidirectional, and there could be multiple valid connections between the same two houses with different costs.

Return the minimum total cost to supply water to all houses.

• new solution update 2022.03.08

class Solution {
    vector<int> g;
    int find(int n) {
        return g[n] == n ? n : g[n] = find(g[n]);
    }
    void uni(int a, int b) {
        int pa = find(a), pb = find(b);
        g[pa] = g[pb] = min(pa,pb);
    }
public:
    int minCostToSupplyWater(int n, vector<int>& wells, vector<vector<int>>& pipes) {
        g = vector<int>(n);
        int conn = 0;
        priority_queue<array<int,3>, vector<array<int,3>>, greater<array<int,3>>> pq;
        for(int i = 0; i < n; i++) g[i] = i;
        for(auto p : pipes) {
            pq.push({p[2], p[0] - 1, p[1] - 1});
        }
        while(!pq.empty()) {
            auto [c, f, t] = pq.top(); pq.pop();
            if(c <= max(wells[find(f)], wells[find(t)]) and find(f) != find(t)) { //merge
                int pa = find(f), pb = find(t);
                wells[min(pa,pb)] = min(wells[pa],wells[pb]);
                uni(f,t);
                conn += c;
            }
        }
        int pump = 0;
        unordered_set<int> s;
        for(int i = 0; i < n; i++) {
            if(s.insert(find(i)).second) {
                pump += wells[find(i)];
            }
        }
        return pump + conn;
    }
};

• Time : O((n+m)logn) [mlogm on heap, mlogn on union find]
• Space : O(n)

class Solution {
    int find(vector<int>& group, int n) {
        return group[n] == n? n : group[n] = find(group, group[n]);
    }
    
    bool uni(vector<int>& group, vector<int>& wells, int a, int b, int cost) {
        int mia = find(group,a), mib = find(group,b);
        if(mia == mib) return false;
        int ca = wells[mia-1], cb = wells[mib-1];
        if(cost > max(ca, cb)) return false;
        if(ca > cb) {
            group[mia] = group[mib] = mib;
        } else if(ca < cb) {
            group[mia] = group[mib] = mia;
        } else {
            group[mia] = group[mib] = min(mia,mib);
        }
        return true;
    }
public:
    int minCostToSupplyWater(int n, vector<int>& wells, vector<vector<int>>& pipes) {
        priority_queue<array<int,3>,vector<array<int,3>>,greater<array<int,3>>> pq;
        for(auto& p : pipes) {
            pq.push({p[2],p[1],p[0]});
        }
        vector<int> group(n+1);
        for(int i = 0; i <= n; i++) group[i] = i;
        int cost = 0;
        
        while(!pq.empty()){
            auto [pCost, a, b] = pq.top(); pq.pop();
            if(uni(group, wells, a, b, pCost)) {
                cost += pCost;
            }
        }
        
        unordered_map<int, int> selectedWell;
        for(int i = 1; i <= n; i++) {
            int mig = find(group,i);
            selectedWell[mig] = wells[mig-1];
        }
        for(auto [gr, well]: selectedWell) {
            cost += well;
        }
        return cost;
    }
};