Advertisement
adityasuman100

Untitled

Mar 1st, 2025
208
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
C++ 5.14 KB | None | 0 0
  1. #include <iostream>
  2. #include <vector>
  3. #include <unordered_set>
  4. using namespace std;
  5.  
  6. class DisjointSet {
  7. private:
  8.     vector<int> parent;
  9.     vector<int> rank;
  10.     vector<unordered_set<int>> group_members;
  11.  
  12. public:
  13.     DisjointSet(int n) {
  14.         parent.resize(n);
  15.         rank.resize(n, 0);
  16.         group_members.resize(n);
  17.        
  18.         for (int i = 0; i < n; i++) {
  19.             parent[i] = i;
  20.             group_members[i].insert(i);
  21.         }
  22.     }
  23.    
  24.     int find(int x) {
  25.         if (parent[x] != x) {
  26.             parent[x] = find(parent[x]);
  27.         }
  28.         return parent[x];
  29.     }
  30.    
  31.     void unite(int x, int y) {
  32.         int rootX = find(x);
  33.         int rootY = find(y);
  34.        
  35.         if (rootX == rootY) return;
  36.        
  37.         if (rank[rootX] < rank[rootY]) {
  38.             parent[rootX] = rootY;
  39.            
  40.             // Move all members from group rootX to group rootY
  41.             for (int member : group_members[rootX]) {
  42.                 group_members[rootY].insert(member);
  43.             }
  44.             group_members[rootX].clear();
  45.         } else {
  46.             parent[rootY] = rootX;
  47.            
  48.             // Move all members from group rootY to group rootX
  49.             for (int member : group_members[rootY]) {
  50.                 group_members[rootX].insert(member);
  51.             }
  52.             group_members[rootY].clear();
  53.            
  54.             if (rank[rootX] == rank[rootY]) {
  55.                 rank[rootX]++;
  56.             }
  57.         }
  58.     }
  59.    
  60.     const unordered_set<int>& getGroupMembers(int x) {
  61.         int root = find(x);
  62.         return group_members[root];
  63.     }
  64. };
  65.  
  66. vector<int> timeOfInfection(int n, int m, vector<int>& initially_infected, vector<vector<int>>& updates) {
  67.     vector<int> infection_time(n, -1);
  68.     vector<bool> vaccinated(n, false);
  69.    
  70.     // Mark initially infected people
  71.     for (int person : initially_infected) {
  72.         infection_time[person] = 0;
  73.     }
  74.    
  75.     DisjointSet ds(n);
  76.    
  77.     // Process each update
  78.     for (int i = 0; i < updates.size(); i++) {
  79.         int time = i + 1;
  80.         int type = updates[i][0];
  81.         int a = updates[i][1];
  82.        
  83.         if (type == 0) {  // Contact between a and b
  84.             int b = updates[i][2];
  85.            
  86.             // Get the current group status before merging
  87.             auto group_a = ds.getGroupMembers(a);
  88.             auto group_b = ds.getGroupMembers(b);
  89.            
  90.             // Check if any infected people in groups
  91.             bool a_group_has_infected = false;
  92.             for (int person : group_a) {
  93.                 if (infection_time[person] != -1) {
  94.                     a_group_has_infected = true;
  95.                     break;
  96.                 }
  97.             }
  98.            
  99.             bool b_group_has_infected = false;
  100.             for (int person : group_b) {
  101.                 if (infection_time[person] != -1) {
  102.                     b_group_has_infected = true;
  103.                     break;
  104.                 }
  105.             }
  106.            
  107.             // Merge the groups
  108.             ds.unite(a, b);
  109.            
  110.             // If one group has infected and the other doesn't,
  111.             // spread infection to unvaccinated members of the other group
  112.             if (a_group_has_infected || b_group_has_infected) {
  113.                 auto merged_group = ds.getGroupMembers(a);  // or ds.getGroupMembers(b), they're the same now
  114.                
  115.                 for (int person : merged_group) {
  116.                     if (infection_time[person] == -1 && !vaccinated[person]) {
  117.                         infection_time[person] = time;
  118.                     }
  119.                 }
  120.             }
  121.         } else {  // Vaccination of person a
  122.             vaccinated[a] = true;
  123.         }
  124.     }
  125.    
  126.     return infection_time;
  127. }
  128.  
  129. int main() {
  130.     // Example 1
  131.     int n1 = 6;
  132.     int m1 = 2;
  133.     vector<int> initially_infected1 = {0, 3};
  134.     vector<vector<int>> updates1 = {
  135.         {0, 0, 3}, {0, 0, 4}, {0, 5, 2},
  136.         {1, 3, -1}, {1, 2, -1}, {0, 2, 4}
  137.     };
  138.    
  139.     vector<int> result1 = timeOfInfection(n1, m1, initially_infected1, updates1);
  140.    
  141.     cout << "Example 1 output: ";
  142.     for (int t : result1) {
  143.         cout << t << " ";
  144.     }
  145.     cout << endl;
  146.    
  147.     // Example 2
  148.     int n2 = 5;
  149.     int m2 = 2;
  150.     vector<int> initially_infected2 = {3, 4};
  151.     vector<vector<int>> updates2 = {
  152.         {1, 1, -1}, {0, 0, 1}, {0, 0, 4}
  153.     };
  154.    
  155.     vector<int> result2 = timeOfInfection(n2, m2, initially_infected2, updates2);
  156.    
  157.     cout << "Example 2 output: ";
  158.     for (int t : result2) {
  159.         cout << t << " ";
  160.     }
  161.     cout << endl;
  162.    
  163.     // Example 3
  164.     int n3 = 6;
  165.     int m3 = 5;
  166.     vector<int> initially_infected3 = {0, 1, 2, 3, 4};
  167.     vector<vector<int>> updates3 = {
  168.         {0, 0, 1}, {0, 0, 2}, {0, 0, 3}, {0, 0, 4}, {0, 1, 2},
  169.         {0, 1, 3}, {0, 1, 4}, {0, 2, 3}, {0, 2, 4}, {0, 3, 4}
  170.     };
  171.    
  172.     vector<int> result3 = timeOfInfection(n3, m3, initially_infected3, updates3);
  173.    
  174.     cout << "Example 3 output: ";
  175.     for (int t : result3) {
  176.         cout << t << " ";
  177.     }
  178.     cout << endl;
  179.    
  180.     return 0;
  181. }
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement