Leetcode 2092 - Find All People With Secret

Link to problem

Problem

You are given an integer n indicating there are n people numbered from 0 to n - 1. You are also given a 0-indexed 2D integer array meetings where meetings[i] = [x_i, y_i, time_i] indicates that person x_i and person yi have a meeting at timei. A person may attend multiple meetings at the same time. Finally, you are given an integer firstPerson.

Person 0 has a secret and initially shares the secret with a person firstPerson at time 0. This secret is then shared every time a meeting takes place with a person that has the secret. More formally, for every meeting, if a person x_i has the secret at time_i, then they will share the secret with person y_i, and vice versa.

The secrets are shared instantaneously. That is, a person may receive the secret and share it with people in other meetings within the same time frame.

Return a list of all the people that have the secret after all the meetings have taken place. You may return the answer in any order.

Solution

If a meeting happened with a person that has a secret, the two participant should be in the same union. We could achieve this fast by using Disjoint Set data structure. See here for some details about disjoint set.

Since the time of the meeting happening is important, we will have to sort all the meetings before processing. Moreover, after each day of meetings, we have to check if the union are the union that shared the secret (by checking if find_set(0) and find_set(i) are the same as node 0 always possess the secret). If the union of nodes does not have the secret, we have to “reset” the union. If reset operations are not done, the union that does not have the secret might be wrongly joined with the union that has the secret in the meetings, and also pulled a large set of nodes that have not meet with any node that possess the secret.

Implementation

class DisjointSet {
public:
    DisjointSet(int n): parent(n) {
        for (int i = 0; i < n; i++) parent[i] = i;
    }

    void union_set(int a, int b) {
        a = find_set(a);
        b = find_set(b);
        if (a != b) parent[b] = a;
    }

    int find_set(int v) {
        if (v == parent[v]) return v;
        return parent[v] = find_set(parent[v]);
    }

    bool is_union(int a, int b) {
        a = find_set(a);
        b = find_set(b);
        return a == b;
    }

    // "reset" the union
    void reset(int v) {
        parent[v] = v;
    }

private:
    vector<int> parent;
};

class Solution {
public:

    vector<int> findAllPeople(int n, vector<vector<int>>& meetings, int firstPerson) {
        vector<int> output;
        DisjointSet dset(n);

        // sort meetings by date
        sort(meetings.begin(), meetings.end(), [](const auto& a, const auto& b) -> bool {
                return a[2] < b[2];
            });

        dset.union_set(0, firstPerson);

        for (int i = 0; i < meetings.size();) {
            vector<int> meeted;
            int cur_date = meetings[i][2];
            for (; i < meetings.size() && meetings[i][2] == cur_date; i++) {
                dset.union_set(meetings[i][0], meetings[i][1]);
                meeted.push_back(meetings[i][0]);
                meeted.push_back(meetings[i][1]);
            }

            for (auto j: meeted) {
                if (!dset.is_union(0, j)) dset.reset(j);
            }
        }

        for (int i = 0; i < n; i++) {
            if (dset.is_union(0, i)) output.push_back(i);
        }

        return output;
    }
};