Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- #include <iostream>
- #include <vector>
- using namespace std;
- class DisjointSet {
- private:
- vector<int> parent;
- vector<int> rank;
- vector<bool> has_infected; // Whether the group has any infected member
- public:
- DisjointSet(int n) {
- parent.resize(n);
- rank.resize(n, 0);
- has_infected.resize(n, false);
- for (int i = 0; i < n; i++) {
- parent[i] = i;
- }
- }
- int find(int x) {
- if (parent[x] != x) {
- parent[x] = find(parent[x]);
- }
- return parent[x];
- }
- void unite(int x, int y) {
- int rootX = find(x);
- int rootY = find(y);
- if (rootX == rootY) return;
- if (rank[rootX] < rank[rootY]) {
- parent[rootX] = rootY;
- has_infected[rootY] = has_infected[rootY] || has_infected[rootX];
- } else {
- parent[rootY] = rootX;
- has_infected[rootX] = has_infected[rootX] || has_infected[rootY];
- if (rank[rootX] == rank[rootY]) {
- rank[rootX]++;
- }
- }
- }
- bool groupHasInfected(int root) {
- return has_infected[root];
- }
- void markGroupAsInfected(int x) {
- int root = find(x);
- has_infected[root] = true;
- }
- };
- vector<int> timeOfInfection(int n, int m, vector<int>& initially_infected, vector<vector<int>>& updates) {
- vector<int> infection_time(n, -1);
- vector<bool> vaccinated(n, false);
- DisjointSet ds(n);
- // Mark initially infected people
- for (int person : initially_infected) {
- infection_time[person] = 0;
- ds.markGroupAsInfected(person);
- }
- // Process each update
- for (int i = 0; i < updates.size(); i++) {
- int time = i + 1;
- int type = updates[i][0];
- int a = updates[i][1];
- if (type == 0) { // Contact between a and b
- int b = updates[i][2];
- int root_a = ds.find(a);
- int root_b = ds.find(b);
- // Check if the groups are already merged
- if (root_a == root_b) continue;
- bool a_group_has_infected = ds.groupHasInfected(root_a);
- bool b_group_has_infected = ds.groupHasInfected(root_b);
- // Track people in group b before merging
- vector<int> group_b_members;
- if (a_group_has_infected && !b_group_has_infected) {
- for (int person = 0; person < n; person++) {
- if (ds.find(person) == root_b) {
- group_b_members.push_back(person);
- }
- }
- }
- // Track people in group a before merging
- vector<int> group_a_members;
- if (!a_group_has_infected && b_group_has_infected) {
- for (int person = 0; person < n; person++) {
- if (ds.find(person) == root_a) {
- group_a_members.push_back(person);
- }
- }
- }
- // Merge the groups
- ds.unite(a, b);
- // Infect unvaccinated members of group b if group a has infected
- if (a_group_has_infected && !b_group_has_infected) {
- for (int person : group_b_members) {
- if (infection_time[person] == -1 && !vaccinated[person]) {
- infection_time[person] = time;
- }
- }
- }
- // Infect unvaccinated members of group a if group b has infected
- if (!a_group_has_infected && b_group_has_infected) {
- for (int person : group_a_members) {
- if (infection_time[person] == -1 && !vaccinated[person]) {
- infection_time[person] = time;
- }
- }
- }
- } else { // Vaccination of person a
- vaccinated[a] = true;
- }
- }
- return infection_time;
- }
- int main() {
- // Example 1
- int n1 = 6;
- int m1 = 2;
- vector<int> initially_infected1 = {0, 3};
- vector<vector<int>> updates1 = {
- {0, 0, 3}, {0, 0, 4}, {0, 5, 2},
- {1, 3, -1}, {1, 2, -1}, {0, 2, 4}
- };
- vector<int> result1 = timeOfInfection(n1, m1, initially_infected1, updates1);
- cout << "Example 1 output: ";
- for (int t : result1) {
- cout << t << " ";
- }
- cout << endl;
- // Example 2
- int n2 = 5;
- int m2 = 2;
- vector<int> initially_infected2 = {3, 4};
- vector<vector<int>> updates2 = {
- {1, 1, -1}, {0, 0, 1}, {0, 0, 4}
- };
- vector<int> result2 = timeOfInfection(n2, m2, initially_infected2, updates2);
- cout << "Example 2 output: ";
- for (int t : result2) {
- cout << t << " ";
- }
- cout << endl;
- // Example 3
- int n3 = 6;
- int m3 = 5;
- vector<int> initially_infected3 = {0, 1, 2, 3, 4};
- vector<vector<int>> updates3 = {
- {0, 0, 1}, {0, 0, 2}, {0, 0, 3}, {0, 0, 4}, {0, 1, 2},
- {0, 1, 3}, {0, 1, 4}, {0, 2, 3}, {0, 2, 4}, {0, 3, 4}
- };
- vector<int> result3 = timeOfInfection(n3, m3, initially_infected3, updates3);
- cout << "Example 3 output: ";
- for (int t : result3) {
- cout << t << " ";
- }
- cout << endl;
- return 0;
- }
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement