#include "game.h"
#include <algorithm>
#include <cassert>
#include <cstddef>
#include <iostream>
#include <map>
#include <optional>
#include <utility>
#include <vector>
using namespace std;
class CityForest {
public:
explicit CityForest(size_t n) {
assert(n > 0);
count = n;
for (size_t i = 0; i < n; ++i) {
parents.push_back(i);
}
for (size_t i = 0; i < n; ++i) {
map<size_t, size_t> i_can_connect_to;
for (size_t j = 0; j < n; ++j) {
cout << "adding " << j << " as a connection of " << i << '\n';
if (i != j) {
i_can_connect_to.insert(
{j, i}); // (component can connect to j from i)
}
}
can_connect_to.push_back(i_can_connect_to);
}
sizes = vector<size_t>(n, 1);
}
bool are_connected(size_t a, size_t b) {
return get_representative(a) == get_representative(b);
}
void ensure_freedom(size_t k) {
map<size_t, size_t> connections = can_connect_to.at(get_representative(k));
if (connections.size() == 1) {
auto connection = connections.begin();
size_t from = connection->second;
size_t to = connection->first;
join(from, to); // mandatory join operation
ensure_freedom(to);
}
}
void denied_connecting_edge(size_t a, size_t b) {
cout << "denied " << a << " and " << b << "\n";
assert(!are_connected(a, b));
map<size_t, size_t> &a_connections =
can_connect_to.at(get_representative(a));
map<size_t, size_t> &b_connections =
can_connect_to.at(get_representative(b));
cout << a << " has " << a_connections.size() << " connections.\n";
cout << b << " has " << b_connections.size() << " connections.\n";
a_connections.erase(b);
b_connections.erase(a);
ensure_freedom(a);
ensure_freedom(b);
}
[[nodiscard]] size_t size_of(size_t k) {
return sizes.at(get_representative(k));
}
[[nodiscard]] size_t size() const { return count; }
private:
vector<size_t> parents;
vector<size_t> sizes;
vector<map<size_t, size_t>> can_connect_to;
size_t count = 0;
void join(size_t absorbed_node, size_t absorbing_node) {
if (size_of(absorbed_node) > size_of(absorbing_node)) {
swap(absorbed_node, absorbing_node);
}
if (are_connected(absorbed_node, absorbing_node)) {
return;
}
cout << "joining " << absorbed_node << " & " << absorbing_node << '\n';
size_t new_size = size_of(absorbed_node) + size_of(absorbing_node);
size_t remaining_representative = get_representative(absorbing_node);
size_t absorbed_representative = get_representative(absorbed_node);
parents.at(absorbed_representative) = remaining_representative;
sizes.at(remaining_representative) = new_size;
assert(can_connect_to.at(remaining_representative).size() == 1 ||
can_connect_to.at(absorbed_representative).size() == 1);
assert(
can_connect_to.at(remaining_representative).count(absorbed_node) == 1 &&
can_connect_to.at(absorbed_representative).count(absorbing_node) == 1);
if (can_connect_to.at(remaining_representative).size() == 1) {
can_connect_to.at(remaining_representative) =
can_connect_to.at(absorbed_representative);
can_connect_to.at(remaining_representative).erase(absorbing_node);
} else if (can_connect_to.at(absorbed_representative).size() == 1) {
can_connect_to.at(remaining_representative).erase(absorbed_node);
}
--count;
}
[[nodiscard]] size_t get_representative(size_t k) {
size_t current = k;
size_t parent = parents.at(k);
while (current != parent) {
current = parent;
parent = parents.at(current);
}
size_t representative = current;
current = k;
parent = parents.at(k);
parents.at(k) = representative;
while (current != parent) {
current = parent;
parent = parents.at(current);
parents.at(current) = representative;
}
return representative;
}
};
optional<CityForest> connected_components; // NOLINT
void initialize(int n) { connected_components = CityForest(n); }
int hasEdge(int u, int v) {
if (connected_components.value().are_connected(u, v)) {
cout << u << " and " << v << " are connected.\n";
return 1;
}
connected_components.value().denied_connecting_edge(u, v);
return 0;
}
