LeetCode/cpp/find-all-people-with-secret.cpp

#include <algorithm>
#include <map>
#include <queue>
#include <unordered_map>
#include <unordered_set>
#include <vector>

class Solution {
    void process_timeslot(std::vector<bool> &knows,
                          const std::vector<std::pair<int, int>> &meetings) {
        // Reconstruct the components for the same timeslot
        // and also construct the source of secret
        std::unordered_set<int> sources;
        std::unordered_map<int, std::vector<int>> neighbors;
        for (const auto &[a, b] : meetings) {
            if (knows[a]) {
                sources.insert(a);
            }
            if (knows[b]) {
                sources.insert(b);
            }

            neighbors[a].push_back(b);
            neighbors[b].push_back(a);
        }

        // Constructing queue from iterators is allowed since C++23…
        std::queue<int> q;
        for (const auto &p : sources) {
            q.push(p);
        }

        // Classic BFS
        while (!q.empty()) {
            int u = q.front();
            q.pop();

            for (const auto &neighbor : neighbors[u]) {
                if (knows[neighbor]) {
                    continue;
                }

                knows[neighbor] = true;
                q.push(neighbor);
            }
        }
    }

  public:
    std::vector<int> findAllPeople(int n,
                                   std::vector<std::vector<int>> meetings,
                                   int firstPerson) {
        // Sort the meetings by time
        std::sort(meetings.begin(), meetings.end(),
                  [](const auto &l, const auto &r) { return l[2] < r[2]; });

        // Group meetings first
        std::map<int, std::vector<std::pair<int, int>>> groups;
        for (const auto &meeting : meetings) {
            auto a = meeting[0];
            auto b = meeting[1];
            auto t = meeting[2];

            groups[t].emplace_back(a, b);
        }

        // Keep “knowers” in the bitset
        std::vector<bool> knows(n, false);
        knows[0] = true;
        knows[firstPerson] = true;

        // Go through the groups
        for (const auto &[time, meetings] : groups) {
            process_timeslot(knows, meetings);
        }

        // Reconstruct the list of people that know
        std::vector<int> all_people;
        for (auto i = 0; i < n; ++i) {
            if (knows[i]) {
                all_people.push_back(i);
            }
        }

        return all_people;
    }
};

#ifdef _MF_TEST
#include <gtest/gtest.h>

TEST(examples, no_1) {
    Solution s;
    EXPECT_EQ(
        (s.findAllPeople(6,
                         std::vector{std::vector{1, 2, 5}, std::vector{2, 3, 8},
                                     std::vector{1, 5, 10}},
                         1)),
        (std::vector{0, 1, 2, 3, 5}));
}
TEST(examples, no_2) {
    Solution s;
    EXPECT_EQ(
        (s.findAllPeople(4,
                         std::vector{std::vector{3, 1, 3}, std::vector{1, 2, 2},
                                     std::vector{0, 3, 3}},
                         3)),
        (std::vector{0, 1, 3}));
}
TEST(examples, no_3) {
    Solution s;
    EXPECT_EQ(
        (s.findAllPeople(5,
                         std::vector{std::vector{3, 4, 2}, std::vector{1, 2, 1},
                                     std::vector{2, 3, 1}},
                         1)),
        (std::vector{0, 1, 2, 3, 4}));
}

TEST(wrong_answer, no_1) {
    Solution s;
    EXPECT_EQ(
        (s.findAllPeople(6,
                         std::vector{std::vector{0, 2, 1}, std::vector{1, 3, 1},
                                     std::vector{4, 5, 1}},
                         1)),
        (std::vector{0, 1, 2, 3}));
}

int main(int argc, char **argv) {
    testing::InitGoogleTest(&argc, argv);
    return RUN_ALL_TESTS();
}
#endif