dcd

#include <iostream>
#include <vector>
#include <unordered_set>
#include <functional>

using namespace std;

class Solution {
public:
    vector<int> timeOfInfection(int n, int m, vector<int>& initially_infected, vector<vector<int>>& updates) {
        vector<int> parent(n), infectionTime(n, -1), vaccinated(n, 0);

        // Initialize Union-Find structure
        for (int i = 0; i < n; ++i) parent[i] = i;

        // Set initial infection times
        for (int person : initially_infected) infectionTime[person] = 0;

        // Define find function with an explicit return type
        function<int(int)> find = [&](int x) {
            if (x != parent[x]) parent[x] = find(parent[x]);
            return parent[x];
        };

        // Helper function to infect all connected components recursively
        function<void(int, int)> propagateInfection = [&](int node, int time) {
            int root = find(node);
            for (int i = 0; i < n; ++i) {
                if (find(i) == root && infectionTime[i] == -1 && !vaccinated[i]) {
                    infectionTime[i] = time; // Set infection time for non-vaccinated individuals
                }
            }
        };

        auto unionGroups = [&](int a, int b, int time) {
            int rootA = find(a);
            int rootB = find(b);
            if (rootA != rootB) {
                // Merge rootB into rootA
                parent[rootB] = rootA;

                if (infectionTime[rootA] != -1 && infectionTime[rootB] == -1) {
                    propagateInfection(rootB, time);
                } else if (infectionTime[rootB] != -1
                && infectionTime[rootA] == -1) {
                    propagateInfection(rootA, time);
                }
            }
        };

        // Process updates
        for (int t = 0; t < updates.size(); ++t) {
            int type = updates[t][0];
            int a = updates[t][1];
            int b = updates[t][2];

            if (type == 0) { // Contact
                unionGroups(a, b, t + 1);
            } else if (type == 1) { // Vaccination
                // Only vaccinate if the person is not infected
                if (infectionTime[a] == -1) {
                    vaccinated[a] = 1;
                }
            }
        }

        return infectionTime;
    }
};

int main() {
    // Sample Input
    int n = 5; // Number of people
    int m = 2; // Number of initially infected people
    vector<int> initially_infected = {1, 4}; // Initially infected people
    vector<vector<int>> update = {
        {0, 1, 2},
        {1, 0, -1},
        {0, 0, 2},
        {0, 0, 3}
    };

    // Solution
    Solution sol;
    vector<int> result = sol.timeOfInfection(n, m, initially_infected, update);

    // Output the result
    cout << "Infection times for each person:" << endl;
    for (int i = 0; i < result.size(); ++i) {
        cout << "Person " << i << ": " << result[i] << endl;
    }

    return 0;
}