Short C++ Union Find Solution with Detailed Explanation
This question looks exactly like 924. Minimize Malware Spread, but this question requires us "completely removing it and any connections from this node to any o
This question looks exactly like [924. Minimize Malware Spread](https://leetcode.com/problems/minimize-malware-spread/), but this question requires us **"completely removing it and any connections from this node to any other node"**\n\nMy solution to 924 using Union Find: [C++ Union Find (Short & Easy to Understand)](https://leetcode.com/problems/minimize-malware-spread/discuss/791608/C%2B%2B-Union-Find-(Short-and-Easy-to-Understand))\n\nThe key steps / ideas in union find solution for this question:\n1. Do union on non-malware nodes into groups\n2. Find which groups are infected by one malware; and how many malware sources are in each group\n3. Try removing each malware, find out how many nodes removing this malware could we save from being infected\n\n**Detailed Steps with Code**:\n\nSimilar to other union find problems, initialize union find data structure `roots` and define `find` function (I used a lambda function)\n```cpp\n\n// Union Find set-up\nint n = graph.size();\nvector<int> roots(n);\nfor (auto i = 0; i < n; i++) roots[i] = i;\n\nauto find = [&](int id) {\n\twhile (roots[id] != id) {\n\t\troots[id] = roots[roots[id]];\n\t\tid = roots[id];\n\t}\n\treturn id;\n};\n```\n\nThen, we find non-malware nodes, i.e. nodes doesn\'t appear in `initial`\n```cpp\n// Consider non-malware nodes\nvector<int> clean;\nset<int> initial_set(initial.begin(), initial.end());\nfor (auto i = 0; i < n; i++)\n\tif (!initial_set.count(i)) clean.push_back(i);\n```\n\nFor these non-malware nodes, we do union\n```cpp\n// Union Find on non-malware nodes\nfor (auto i : clean)\n\tfor (auto j : clean)\n\t\tif (i != j && graph[i][j]) roots[find(i)] = find(j);\n```\n\nFor each group of nodes, we found its size and store it in `area`\n```cpp\n// Size of each group\nvector<int> area(n, 0);\nfor (auto i : clean) area[find(i)]++;\n```\n\nOnce we finished the above steps, we are ready to **add malware** into these nodes\nFor each malware, we try to find which groups it\'s connected to. And for each of these groups, we count the number of malware infecting it (or malware sources).\n```cpp\nmap<int, set<int>> infect_node; // malware -> nodes infected\nmap<int, int> infect_count; // node -> # of infecting malwares\nfor (auto i : initial) {\n\tfor (auto j : clean) {\n\t\tif (graph[i][j]) infect_node[i].insert(find(j));\n\t}\n\tfor (auto j : infect_node[i]) infect_count[j] += 1;\n}\n```\n\nThe final step. We try to remove each malware (and all the groups it\'s connecting). If the group has exactly one malware source, then removing this malware should "save" this group of nodes from malware. Otherwise, if the group has multiple malware sources, removing this malware doesn\'t help because the nodes will eventually infected by another malware.  \nBy doing this, we could find out the maximum nodes saved by removing a malware node.\n```cpp\n// Try remove each malware\nint res = initial[0], max_cnt = -1;\nfor (auto& [malware, nodes] : infect_node) {\n\tint count = 0;\n\tfor (auto& node : nodes) {\n\t\t// Only effective when node is infected by one malware\n\t\tif (infect_count[node] == 1)\n\t\t\tcount += area[node];\n\t}\n\tif (count > max_cnt || (count >= max_cnt && malware < res)) {\n\t\tmax_cnt = count;\n\t\tres = malware;\n\t}\n}\n```\n\n\n**Full solution**:\n```cpp\nclass Solution {\npublic:\n    int minMalwareSpread(vector<vector<int>>& graph, vector<int>& initial) {\n        int n = graph.size();\n        vector<int> roots(n);\n        for (auto i = 0; i < n; i++) roots[i] = i;\n        \n        auto find = [&](int id) {\n            while (roots[id] != id) {\n                roots[id] = roots[roots[id]];\n                id = roots[id];\n            }\n            return id;\n        };\n        \n        // Consider non-malware nodes\n        vector<int> clean;\n        set<int> initial_set(initial.begin(), initial.end());\n        for (auto i = 0; i < n; i++)\n            if (!initial_set.count(i)) clean.push_back(i);\n        \n        // Union Find on non-malware nodes\n        for (auto i : clean)\n            for (auto j : clean)\n                if (i != j && graph[i][j])\n                    roots[find(i)] = find(j);\n        \n        // Size of each group\n        vector<int> area(n, 0);\n        for (auto i : clean) area[find(i)]++;\n        \n        map<int, set<int>> infect_node; // malware -> nodes infected\n        map<int, int> infect_count; // node -> # of infecting malware\n        for (auto i : initial) {\n            for (auto j : clean) {\n                if (graph[i][j]) infect_node[i].insert(find(j));\n            }\n            for (auto j : infect_node[i])\n                infect_count[j] += 1;\n        }\n        \n        // Try remove each malware\n        int res = initial[0], max_cnt = -1;\n        for (auto& [malware, nodes] : infect_node) {\n            int count = 0;\n            for (auto& node : nodes) {\n                // Only effective when node is infected by one malware\n                if (infect_count[node] == 1)\n                    count += area[node];\n            }\n            if (count > max_cnt || (count >= max_cnt && malware < res)) {\n                max_cnt = count;\n                res = malware;\n            }\n        }\n        \n        return res;\n    }\n};\n```\n\n\nCredit:\n- https://leetcode.com/problems/minimize-malware-spread-ii/discuss/253064/python-union-find-beat-100
[]