#include "islands.h"
#include <iostream>
#include <map>
#include <queue>
#include <variant>
#include <vector>
std::variant<bool, std::vector<int>> find_journey_connected(int N, int M, std::vector<int> U, std::vector<int> V, std::vector<int> original_labels);
std::variant<bool, std::vector<int>> find_journey(int N, int M, std::vector<int> U, std::vector<int> V) {
std::vector<std::vector<int>> canoesfrom(N);
for (int i = 0; i < M; ++i) canoesfrom[U[i]].push_back(i);
std::vector<bool> seen(N);
std::vector<int> stack = {0};
while (!stack.empty()) {
int node = stack.back();
stack.pop_back();
if (seen[node]) continue;
seen[node] = true;
for (int i : canoesfrom[node]) {
stack.push_back(V[i]);
}
}
std::vector<int> newu, newv, original_labels;
int newm = 0;
for (int i = 0; i < N; ++i) {
if (!seen[i]) continue;
for (int out : canoesfrom[i]) {
newu.push_back(U[out]);
newv.push_back(V[out]);
original_labels.push_back(out);
++newm;
}
}
return find_journey_connected(N, newm, newu, newv, original_labels);
}
std::variant<bool, std::vector<int>> find_journey_connected(int N, int M, std::vector<int> U, std::vector<int> V, std::vector<int> original_labels) {
std::vector<std::vector<int>> out(N), in(N);
for (int i = 0; i < M; ++i) out[U[i]].push_back(i), in[V[i]].push_back(i);
U.push_back(-1);
V.push_back(0); // dummy edge to source
// stage 1: find any loop
std::vector<int> is_marked(N);
std::vector<int> seen(N);
std::vector<int> inpath(N);
std::vector<int> path;
std::vector<int> cycle;
std::vector<int> stack = {M};
bool foundcyclestagei = false;
while (!stack.empty()) {
int canoe = stack.back();
stack.pop_back();
if (canoe == -1) {
inpath[V[path.back()]] = false;
path.pop_back();
continue;
}
path.push_back(canoe);
int node = V[canoe];
if (seen[node]) {
if (!inpath[node]) continue;
foundcyclestagei = true;
is_marked[node] = true;
while (U[path.back()] != node) {
cycle.push_back(path.back());
is_marked[U[path.back()]] = true;
path.pop_back();
}
cycle.push_back(path.back());
path.pop_back();
break;
}
seen[node] = true;
inpath[node] = true;
for (int to : out[node]) {
stack.push_back(-1);
stack.push_back(to);
}
}
if (!foundcyclestagei) return false;
// Stage 2: check multiple reachability to marked nodes
std::vector<long long> weights(M + 1);
for (int canoe : path) {
weights[canoe] = 1;
}
for (int canoe : cycle) {
weights[canoe] = 1e9;
}
std::vector<int> distance(N, 1e9);
std::priority_queue<std::pair<int, int>, std::vector<std::pair<int, int>>, std::greater<std::pair<int, int>>> q;
q.emplace(0, 0);
bool multiplereachabilitystageii = false;
int max_weight = path.size() - 1;
while (!q.empty()) {
auto [d, node] = q.top();
q.pop();
if (d >= distance[node]) continue;
distance[node] = d;
if (is_marked[node] && d <= max_weight - (node == V[path.back()])) multiplereachabilitystageii = true;
for (int to : out[node]) {
q.emplace(d + weights[to], V[to]);
}
}
if (multiplereachabilitystageii) return true;
// Stage 3: we don't have multiple reachability, but let's look for a second cycle
bool foundcyclestageiii = false;
stack = {M};
seen = std::vector<int>(N, false);
while (!stack.empty()) {
int canoe = stack.back();
stack.pop_back();
int node = V[canoe];
if (is_marked[node] && node != path.back()) continue; // avoid refinding the same cycle
if (seen[node]) {
foundcyclestageiii = true;
break;
}
seen[node] = true;
for (int to : out[node]) {
stack.push_back(to);
}
}
return foundcyclestageiii;
}
