infection 2nd sol

#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 x) {
        int root = find(x);
        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);

            bool a_group_has_infected = ds.groupHasInfected(a);
            bool b_group_has_infected = ds.groupHasInfected(b);

            // Merge the groups
            ds.unite(a, b);

            // If one group has infected members, we need to spread infection
            if (a_group_has_infected || b_group_has_infected) {
                // We need a second pass through all people to check who gets infected
                // This is more efficient than storing sets of group members
                for (int person = 0; person < n; person++) {
                    if (ds.find(person) == ds.find(a) && 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;
}