Minimize Malware Spread Solutions in C++

/// Source : https://leetcode.com/problems/minimize-malware-spread/description/
/// Author : liuyubobobo
/// Time   : 2018-10-13
/// Updated: 2019-09-23

#include <iostream>
#include <vector>
#include <set>

using namespace std;


/// DFS
/// Time Complexity: O(n^2)
/// Space Complexity: O(n)
class Solution {

private:
    int n;
    set<int> init_set;
    vector<bool> visited;
    vector<int> cc_sz;
    vector<int> cc_mal_num;

public:
    int minMalwareSpread(vector<vector<int>>& graph, vector<int>& initial) {

        for(int e: initial)
            init_set.insert(e);

        n = graph.size();
        visited.resize(n, false);

        int best = 0, res = -1, ccid = 0;
        for(int v: init_set)
            if(!visited[v]){
                cc_sz.push_back(0), cc_mal_num.push_back(0);
                cc_sz[ccid] = dfs(graph, v, ccid);
                if(cc_mal_num[ccid] == 1 && cc_sz[ccid] > best){
                    best = cc_sz[ccid];
                    res = v;
                }
                ccid ++;
            }
        return res == -1 ? *init_set.begin() : res;
    }

private:
    int dfs(const vector<vector<int>>& g, int v, int ccid){

        visited[v] = true;
        if(init_set.count(v)) cc_mal_num[ccid] ++;

        int res = 1;
        for(int next = 0; next < n; next ++)
            if(g[v][next] && !visited[next])
                res += dfs(g, next, ccid);
        return res;
    }
};


int main() {

    return 0;
}

/// Source : https://leetcode.com/problems/minimize-malware-spread/description/
/// Author : liuyubobobo
/// Time   : 2018-10-16
/// Updated: 2019-09-23

#include <iostream>
#include <vector>
#include <unordered_map>
#include <set>

using namespace std;


/// Using Union-Find
/// Time Complexity: O(n^2)
/// Space Complexity: O(n)
class UF{

private:
    vector<int> parent, sz;

public:
    UF(int n){
        parent.clear();
        for(int i = 0 ; i < n ; i ++){
            parent.push_back(i);
            sz.push_back(1);
        }
    }

    int find(int p){
        if( p != parent[p] )
            parent[p] = find( parent[p] );
        return parent[p];
    }

    bool isConnected(int p , int q){
        return find(p) == find(q);
    }

    void unionElements(int p, int q){

        int pRoot = find(p);
        int qRoot = find(q);

        if( pRoot == qRoot )
            return;

        parent[pRoot] = qRoot;
        sz[qRoot] += sz[pRoot];
    }

    int size(int p){
        return sz[find(p)];
    }
};

class Solution {

public:
    int minMalwareSpread(vector<vector<int>>& graph, vector<int>& initial) {

        int n = graph.size();
        UF uf(n);
        for(int i = 0; i < n; i ++)
            for(int j = i + 1; j < n; j ++)
                if(graph[i][j])
                    uf.unionElements(i, j);

        set<int> init_set(initial.begin(), initial.end());

        unordered_map<int, int> cc_sz, cc_mal_num;
        for(int e: init_set)
            cc_sz[uf.find(e)] = uf.size(e),
            cc_mal_num[uf.find(e)] ++;

        int best = 0, res = -1;
        for(int e: init_set)
            if(cc_mal_num[uf.find(e)] == 1 && cc_sz[uf.find(e)] > best){
                best = cc_sz[uf.find(e)];
                res = e;
            }
        return res == -1 ? *init_set.begin() : res;
    }
};


int main() {

    return 0;
}