Leetcode 1168: Optimize Water Distribution in a Village

https://leetcode.com/problems/optimize-water-distribution-in-a-village/description/

Notes:

This is another question about minimum spanning tree (MST). But we need to convert "digging wells" into "connections": for example, connections to well 0. Thus, it is become the same as Leetcode 1135: Connecting Cities With Minimum Cost 

The key point is it will use the standard union-find method. And the algorithm for MST is Kruskal's algorithm which is in a greedy fashion.

See the code below:

class Solution {
public:
    int minCostToSupplyWater(int n, vector<int>& wells, vector<vector<int>>& pipes) {
        int res = 0;
        for(int i=1; i<=n; ++i) pipes.push_back({0, i, wells[i-1]});
        sort(pipes.begin(), pipes.end(), [](auto &a, auto &b) {return a[2] < b[2];});
        vector<int> f(n+1, 0);
        for(int i=0; i<f.size(); ++i) f[i] = i;
        int ct = 0;
        for(auto &p : pipes) {
            int a = find(f, p[0]), b = find(f, p[1]);
            if(a != b) {
                f[b] = a;
                res += p[2];
                ++ct;
                if(ct==n) break;
            }
        }
        return res;
    }
private:
    int find(vector<int> &f, int x) {
        if(f[x] != x) f[x] = find(f, f[x]);
        return f[x];
    }
};

Here is another implementation using object-oriented design model. See the code below:

class unite_find {
private:
    vector<int> rt;
public:
    unite_find(int x) { //constructor
        for(int i=0; i<x; ++i) rt.push_back(i);
    }
   
    ~unite_find() {};//destructor

    int find(int x) {
        if(rt[x] != x) rt[x] = find(rt[x]);
        return rt[x];
    }

    void unite(int x, int y) {
        rt[x] = y;
    }
};

class Solution {
public:
    int minCostToSupplyWater(int n, vector<int>& wells, vector<vector<int>>& pipes) {
        int res = 0;
        for(int i=1; i<=n; ++i) pipes.push_back({0, i, wells[i-1]});
        sort(pipes.begin(), pipes.end(), [](const auto &a, const auto &b) {return a[2] < b[2];});
        unite_find uf(n+1);
        for(auto &p : pipes) {
            int a = uf.find(p[0]), b = uf.find(p[1]);
            if(a != b) {
                uf.unite(a, b);
                res += p[2];
            }
        }
        return res;
    }
};