#include "train.h"
#include <algorithm>
#include <vector>
const int MAXN = 5e3;
int owner[MAXN], charge[MAXN];
std::vector<int> graph[MAXN];
int degree[MAXN];
std::vector<int> undet;
std::vector<int> removing;
bool toRemove[MAXN], removed[MAXN];
int currentDeg[MAXN];
void reset()
{
for (int node : undet)
{
toRemove[node] = false;
removed[node] = false;
currentDeg[node] = degree[node];
}
}
void flipowners()
{
for (int node : undet)
owner[node] = 1 - owner[node];
}
void removal()
{
for (int node : undet)
if (toRemove[node])
{
removing.emplace_back(node);
removed[node] = true;
}
while (!removing.empty())
{
int rm = removing.back();
removing.pop_back();
for (int next : graph[rm])
if (!removed[next])
{
currentDeg[next]--;
if (owner[next] == 1 || currentDeg[next] == 0)
{
removing.emplace_back(next);
removed[next] = true;
}
}
}
}
void findlosers()
{
reset();
for (int node : undet)
toRemove[node] = charge[node];
removal();
std::vector<int> losers;
for (int node : undet)
if (!removed[node])
losers.emplace_back(node);
reset();
flipowners();
for (int node : losers)
toRemove[node] = true;
removal();
flipowners();
}
std::vector<int> who_wins(std::vector<int> a, std::vector<int> r, std::vector<int> u, std::vector<int> v)
{
int stations = a.size();
int tracks = u.size();
std::vector<int> winner(stations, 0);
for (int iNode = 0; iNode < stations; iNode++)
{
owner[iNode] = a[iNode];
charge[iNode] = r[iNode];
undet.emplace_back(iNode);
}
for (int iEdge = 0; iEdge < tracks; iEdge++)
{
graph[v[iEdge]].emplace_back(u[iEdge]);
degree[u[iEdge]]++;
}
while (!undet.empty())
{
findlosers();
std::vector<int> losers;
for (int node : undet)
if (removed[node])
losers.emplace_back(node);
if (losers.empty())
{
for (int node : undet)
winner[node] = 1;
undet.clear();
}
else
{
for (int node : losers)
{
for (int next : graph[node])
degree[next]--;
undet.erase(std::find(undet.begin(), undet.end(), node));
}
}
}
return winner;
}
