134 lines
3.7 KiB
C++
134 lines
3.7 KiB
C++
#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
|