Leetcode 1579: Remove Max Number of Edges to Keep Graph Fully Traversable

 https://leetcode.com/problems/remove-max-number-of-edges-to-keep-graph-fully-traversable/description/

Notes: 

This question can be solved using union-find. 

The basic idea is that: if there are type 3 edges, we should use it with the highest priority, since it means both players can use it. 

Then the next key point is that: for each edge, we will do a find search in the union-find frame. If the root of the two vertexes are the same, meaning that they have be connected already, then this edge can be removed. Otherwise, it means these two vertexes are not connected yet, so need to use the current edge. 

Another information needs to record is that: we need to whether both players can visited all vertexes or not. This can also be solved in union-find: if there is only one connected group left, then the player can visit all the vertexes. So we just need to count the number of groups in the process of union-find. 

See the code below:

 

class Solution {
public:
    int maxNumEdgesToRemove(int n, vector<vector<int>>& edges) {
        int res = 0;
        vector<pair<int, int>> vt1, vt2, vt3;
        for(auto &a : edges) {
            if(a[0] == 1) vt1.push_back({a[1]-1, a[2]-1});
            else if (a[0] == 2) vt2.push_back({a[1]-1, a[2]-1});
            else vt3.push_back({a[1]-1, a[2]-1});
        }
        vector<int> uf1(n, 0), uf2(n, 0);
        // initialize number of groups (each node is one group initially)
        int ct1 = n, ct2 = n;
        for(int i=1; i<n; ++i) {
            uf1[i] = i;
            uf2[i] = i;
        }
        // use type 3 first
        for(auto &v : vt3) {
            int a = v.first, b = v.second;
            int x = find(a, uf1), y = find(b, uf1);
            // the edge can be revomed;
            if(x==y) ++res;
            // the numer of groups decreases by one
            else { uf1[y] = x; --ct1; uf2[y] = x; --ct2;}
        }
        for(auto &v : vt1) {
            int a = v.first, b = v.second;
            int x = find(a, uf1), y = find(b, uf1);
            if(x==y) ++res;
            else {uf1[y] = x; --ct1;}
        }
        for(auto &v : vt2) {
            int a = v.first, b = v.second;
            int x = find(a, uf2), y = find(b, uf2);
            if(x==y) ++res;
            else {uf2[y] = x; --ct2;}
        }
        // check the final number of groups
        if(ct1>1 || ct2>1) return -1;
        return res;
    }
private:
    int find(int x, vector<int>& vt) {
        if(vt[x] == x) return x;
        return vt[x] = find(vt[x], vt);
    }
};

Another version of code, which is more formal:

 

class Solution {
public:
    struct UF {
        vector<int> ps;
        int num;
        UF(int n): ps(n), num(n) {
            iota(ps.begin(), ps.end(), 0);
        }
        int find(int x) {
            if(ps[x] == x) return x;
            return ps[x] = find(ps[x]);
        }
        int unite(int x, int y) {
            int a = find(x), b = find(y);
            if(a == b) return 1;
            ps[b] = a;
            --num;
            return 0;
        }
        int getNum() {return num;}
    };
    int maxNumEdgesToRemove(int n, vector<vector<int>>& edges) {
        int res = 0;
        vector<pair<int, int>> vt1, vt2, vt3;
        for(auto &a : edges) {
            if(a[0] == 1) vt1.push_back({a[1]-1, a[2]-1});
            else if (a[0] == 2) vt2.push_back({a[1]-1, a[2]-1});
            else vt3.push_back({a[1]-1, a[2]-1});
        }
        UF uf1(n), uf2(n);
        for(auto &v : vt3) {
            res += uf1.unite(v.first, v.second);
            uf2.unite(v.first, v.second);
        }
        for(auto &v : vt1) {
            res += uf1.unite(v.first, v.second);
        }
        for(auto &v : vt2) {
            res += uf2.unite(v.first, v.second);
        }
        if(uf1.getNum()>1 || uf2.getNum()>1) return -1;
        return res;
    }
};