minimizeMalwareSpreadII.cpp

// Source: https://leetcode.com/problems/minimize-malware-spread-ii/
// Author: Miao Zhang
// Date:   2021-03-26

class Solution {
public:
    int minMalwareSpread(vector<vector<int>>& graph, vector<int>& initial) {
        int n = graph.size();
        vector<int> clean(n, 1);
        for (auto& x: initial) {
            clean[x] = 0;
        }
        vector<vector<int>> infected_by(n);
        for (int& x: initial) {
            set<int> seen;
            dfs(graph, clean, x, seen);
            for (int v: seen) {
                infected_by[v].push_back(x);
            }
        }

        vector<int> contribution(n);
        for (int v = 0; v < n; v++) {
            if (infected_by[v].size() == 1) {
                contribution[infected_by[v][0]]++;
            }
        }

        int res = *min_element(begin(initial), end(initial));
        int resscore = -1;
        for (int x: initial) {
            int score = contribution[x];
            if (score > resscore || score == resscore && x < res) {
                res = x;
                resscore = score;
            }
        }
        return res;
    }

private:
    void dfs(vector<vector<int>>& graph, vector<int>& clean, int u, set<int>& seen) {
        for (int v = 0; v < graph.size(); v++) {
            if (graph[u][v] == 1 && clean[v] == 1 && !seen.count(v)) {
                seen.insert(v);
                dfs(graph, clean, v, seen);
            }
        }
    }
};