#include <iostream>
#include <vector>
#include <unordered_set>
#include <algorithm>
using namespace std;
class Solution {
public:
int g_nodes, g_edges;
vector<vector<int>> adj;
// DFS to calculate the number of infected nodes
int dfs(int node, vector<bool>& visited, vector<int>& malware) {
visited[node] = true;
int infectedCount = 1; // Start counting the current infected node
for (int neighbor : adj[node]) {
if (!visited[neighbor] && malware[neighbor] == 1) {
infectedCount += dfs(neighbor, visited, malware);
}
}
return infectedCount;
}
// Main function to remove one node and minimize the malware spread
int minMalwareSpread(vector<int>& g_from, vector<int>& g_to, vector<int>& malware) {
adj.resize(g_nodes);
// Build the graph
for (int i = 0; i < g_edges; ++i) {
adj[g_from[i] - 1].push_back(g_to[i] - 1);
adj[g_to[i] - 1].push_back(g_from[i] - 1);
}
int minInfected = g_nodes + 1; // Initialize with a large number
int result = -1;
// Try removing each infected node
for (int i = 0; i < g_nodes; ++i) {
if (malware[i] == 1) {
// Create a copy of the malware array and remove the current node
vector<int> malwareCopy = malware;
malwareCopy[i] = 0;
// Perform DFS to calculate the spread of malware
vector<bool> visited(g_nodes, false);
int infectedCount = 0;
for (int j = 0; j < g_nodes; ++j) {
if (!visited[j] && malwareCopy[j] == 1) {
infectedCount += dfs(j, visited, malwareCopy);
}
}
// Update the result if this removal minimizes the infected count
if (infectedCount < minInfected) {
minInfected = infectedCount;
result = i;
} else if (infectedCount == minInfected && i < result) {
result = i;
}
}
}
return result;
}
};
int main() {
int g_nodes = 5;
int g_edges = 4;
vector<int> g_from = {1, 2, 3, 4};
vector<int> g_to = {2,3,4,5};
vector<int> malware = {0,0,0,0,0};
Solution sol;
sol.g_nodes = g_nodes;
sol.g_edges = g_edges;
int nodeToRemove = sol.minMalwareSpread(g_from, g_to, malware);
cout << "Remove node: " << nodeToRemove + 1 << endl; // Convert 0-indexed to 1-indexed
return 0;
}
I2luY2x1ZGUgPGlvc3RyZWFtPgojaW5jbHVkZSA8dmVjdG9yPgojaW5jbHVkZSA8dW5vcmRlcmVkX3NldD4KI2luY2x1ZGUgPGFsZ29yaXRobT4KdXNpbmcgbmFtZXNwYWNlIHN0ZDsKCmNsYXNzIFNvbHV0aW9uIHsKcHVibGljOgogICAgaW50IGdfbm9kZXMsIGdfZWRnZXM7CiAgICB2ZWN0b3I8dmVjdG9yPGludD4+IGFkajsKICAgIAogICAgLy8gREZTIHRvIGNhbGN1bGF0ZSB0aGUgbnVtYmVyIG9mIGluZmVjdGVkIG5vZGVzCiAgICBpbnQgZGZzKGludCBub2RlLCB2ZWN0b3I8Ym9vbD4mIHZpc2l0ZWQsIHZlY3RvcjxpbnQ+JiBtYWx3YXJlKSB7CiAgICAgICAgdmlzaXRlZFtub2RlXSA9IHRydWU7CiAgICAgICAgaW50IGluZmVjdGVkQ291bnQgPSAxOyAvLyBTdGFydCBjb3VudGluZyB0aGUgY3VycmVudCBpbmZlY3RlZCBub2RlCiAgICAgICAgCiAgICAgICAgZm9yIChpbnQgbmVpZ2hib3IgOiBhZGpbbm9kZV0pIHsKICAgICAgICAgICAgaWYgKCF2aXNpdGVkW25laWdoYm9yXSAmJiBtYWx3YXJlW25laWdoYm9yXSA9PSAxKSB7CiAgICAgICAgICAgICAgICBpbmZlY3RlZENvdW50ICs9IGRmcyhuZWlnaGJvciwgdmlzaXRlZCwgbWFsd2FyZSk7CiAgICAgICAgICAgIH0KICAgICAgICB9CiAgICAgICAgcmV0dXJuIGluZmVjdGVkQ291bnQ7CiAgICB9CiAgICAKICAgIC8vIE1haW4gZnVuY3Rpb24gdG8gcmVtb3ZlIG9uZSBub2RlIGFuZCBtaW5pbWl6ZSB0aGUgbWFsd2FyZSBzcHJlYWQKICAgIGludCBtaW5NYWx3YXJlU3ByZWFkKHZlY3RvcjxpbnQ+JiBnX2Zyb20sIHZlY3RvcjxpbnQ+JiBnX3RvLCB2ZWN0b3I8aW50PiYgbWFsd2FyZSkgewogICAgICAgIGFkai5yZXNpemUoZ19ub2Rlcyk7CiAgICAgICAgCiAgICAgICAgLy8gQnVpbGQgdGhlIGdyYXBoCiAgICAgICAgZm9yIChpbnQgaSA9IDA7IGkgPCBnX2VkZ2VzOyArK2kpIHsKICAgICAgICAgICAgYWRqW2dfZnJvbVtpXSAtIDFdLnB1c2hfYmFjayhnX3RvW2ldIC0gMSk7CiAgICAgICAgICAgIGFkaltnX3RvW2ldIC0gMV0ucHVzaF9iYWNrKGdfZnJvbVtpXSAtIDEpOwogICAgICAgIH0KICAgICAgICAKICAgICAgICBpbnQgbWluSW5mZWN0ZWQgPSBnX25vZGVzICsgMTsgLy8gSW5pdGlhbGl6ZSB3aXRoIGEgbGFyZ2UgbnVtYmVyCiAgICAgICAgaW50IHJlc3VsdCA9IC0xOwogICAgICAgIAogICAgICAgIC8vIFRyeSByZW1vdmluZyBlYWNoIGluZmVjdGVkIG5vZGUKICAgICAgICBmb3IgKGludCBpID0gMDsgaSA8IGdfbm9kZXM7ICsraSkgewogICAgICAgICAgICBpZiAobWFsd2FyZVtpXSA9PSAxKSB7CiAgICAgICAgICAgICAgICAvLyBDcmVhdGUgYSBjb3B5IG9mIHRoZSBtYWx3YXJlIGFycmF5IGFuZCByZW1vdmUgdGhlIGN1cnJlbnQgbm9kZQogICAgICAgICAgICAgICAgdmVjdG9yPGludD4gbWFsd2FyZUNvcHkgPSBtYWx3YXJlOwogICAgICAgICAgICAgICAgbWFsd2FyZUNvcHlbaV0gPSAwOwogICAgICAgICAgICAgICAgCiAgICAgICAgICAgICAgICAvLyBQZXJmb3JtIERGUyB0byBjYWxjdWxhdGUgdGhlIHNwcmVhZCBvZiBtYWx3YXJlCiAgICAgICAgICAgICAgICB2ZWN0b3I8Ym9vbD4gdmlzaXRlZChnX25vZGVzLCBmYWxzZSk7CiAgICAgICAgICAgICAgICBpbnQgaW5mZWN0ZWRDb3VudCA9IDA7CiAgICAgICAgICAgICAgICAKICAgICAgICAgICAgICAgIGZvciAoaW50IGogPSAwOyBqIDwgZ19ub2RlczsgKytqKSB7CiAgICAgICAgICAgICAgICAgICAgaWYgKCF2aXNpdGVkW2pdICYmIG1hbHdhcmVDb3B5W2pdID09IDEpIHsKICAgICAgICAgICAgICAgICAgICAgICAgaW5mZWN0ZWRDb3VudCArPSBkZnMoaiwgdmlzaXRlZCwgbWFsd2FyZUNvcHkpOwogICAgICAgICAgICAgICAgICAgIH0KICAgICAgICAgICAgICAgIH0KICAgICAgICAgICAgICAgIAogICAgICAgICAgICAgICAgLy8gVXBkYXRlIHRoZSByZXN1bHQgaWYgdGhpcyByZW1vdmFsIG1pbmltaXplcyB0aGUgaW5mZWN0ZWQgY291bnQKICAgICAgICAgICAgICAgIGlmIChpbmZlY3RlZENvdW50IDwgbWluSW5mZWN0ZWQpIHsKICAgICAgICAgICAgICAgICAgICBtaW5JbmZlY3RlZCA9IGluZmVjdGVkQ291bnQ7CiAgICAgICAgICAgICAgICAgICAgcmVzdWx0ID0gaTsKICAgICAgICAgICAgICAgIH0gZWxzZSBpZiAoaW5mZWN0ZWRDb3VudCA9PSBtaW5JbmZlY3RlZCAmJiBpIDwgcmVzdWx0KSB7CiAgICAgICAgICAgICAgICAgICAgcmVzdWx0ID0gaTsKICAgICAgICAgICAgICAgIH0KICAgICAgICAgICAgfQogICAgICAgIH0KICAgICAgICAKICAgICAgICByZXR1cm4gcmVzdWx0OwogICAgfQp9OwoKaW50IG1haW4oKSB7CiAgICBpbnQgZ19ub2RlcyA9IDU7CiAgICBpbnQgZ19lZGdlcyA9IDQ7CiAgICB2ZWN0b3I8aW50PiBnX2Zyb20gPSB7MSwgMiwgMywgNH07CiAgICB2ZWN0b3I8aW50PiBnX3RvID0gezIsMyw0LDV9OwogICAgdmVjdG9yPGludD4gbWFsd2FyZSA9IHswLDAsMCwwLDB9OwogICAgCiAgICBTb2x1dGlvbiBzb2w7CiAgICBzb2wuZ19ub2RlcyA9IGdfbm9kZXM7CiAgICBzb2wuZ19lZGdlcyA9IGdfZWRnZXM7CiAgICBpbnQgbm9kZVRvUmVtb3ZlID0gc29sLm1pbk1hbHdhcmVTcHJlYWQoZ19mcm9tLCBnX3RvLCBtYWx3YXJlKTsKICAgIAogICAgY291dCA8PCAiUmVtb3ZlIG5vZGU6ICIgPDwgbm9kZVRvUmVtb3ZlICsgMSA8PCBlbmRsOyAvLyBDb252ZXJ0IDAtaW5kZXhlZCB0byAxLWluZGV4ZWQKICAgIHJldHVybiAwOwp9Cg==