Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- #include <iostream>
- #include <vector>
- #include <unordered_set>
- using namespace std;
- class DisjointSet {
- private:
- vector<int> parent;
- vector<int> rank;
- vector<unordered_set<int>> group_members;
- public:
- DisjointSet(int n) {
- parent.resize(n);
- rank.resize(n, 0);
- group_members.resize(n);
- for (int i = 0; i < n; i++) {
- parent[i] = i;
- group_members[i].insert(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;
- // Move all members from group rootX to group rootY
- for (int member : group_members[rootX]) {
- group_members[rootY].insert(member);
- }
- group_members[rootX].clear();
- } else {
- parent[rootY] = rootX;
- // Move all members from group rootY to group rootX
- for (int member : group_members[rootY]) {
- group_members[rootX].insert(member);
- }
- group_members[rootY].clear();
- if (rank[rootX] == rank[rootY]) {
- rank[rootX]++;
- }
- }
- }
- const unordered_set<int>& getGroupMembers(int x) {
- int root = find(x);
- return group_members[root];
- }
- };
- 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);
- // Mark initially infected people
- for (int person : initially_infected) {
- infection_time[person] = 0;
- }
- DisjointSet ds(n);
- // 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];
- // Get the current group status before merging
- auto group_a = ds.getGroupMembers(a);
- auto group_b = ds.getGroupMembers(b);
- // Check if any infected people in groups
- bool a_group_has_infected = false;
- for (int person : group_a) {
- if (infection_time[person] != -1) {
- a_group_has_infected = true;
- break;
- }
- }
- bool b_group_has_infected = false;
- for (int person : group_b) {
- if (infection_time[person] != -1) {
- b_group_has_infected = true;
- break;
- }
- }
- // Merge the groups
- ds.unite(a, b);
- // If one group has infected and the other doesn't,
- // spread infection to unvaccinated members of the other group
- if (a_group_has_infected || b_group_has_infected) {
- auto merged_group = ds.getGroupMembers(a); // or ds.getGroupMembers(b), they're the same now
- for (int person : merged_group) {
- 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