#include "train.h"
#include<algorithm>
std::vector<int> who_wins(
std::vector<int> a_own,
std::vector<int> recharge,
std::vector<int> u, std::vector<int> v) {
std::vector<std::vector<int>> ad(a_own.size());
for(unsigned i=0;i<u.size();++i){
ad[u[i]].push_back(v[i]);
}
u.clear();v.clear();
std::vector<std::vector<char>> f(ad.size()+1,std::vector<char>(ad.size()));
// f[k][i] = can B keep the train uncharged from node [i] in [k] steps?
std::transform(begin(recharge),end(recharge),begin(f[0]),[](int i){
return i==0;
});
for(unsigned k=1;k<f.size();++k){
char* fk=f[k].data();
char const* fk_=f[k-1].data();
for(unsigned i=0;i<ad.size();++i){
auto const next_f=[&](int j){return fk_[j];};
if(recharge[i])
fk[i]=false;
else if(a_own[i])
fk[i]=std::all_of(begin(ad[i]),end(ad[i]),next_f);
else
fk[i]=std::any_of(begin(ad[i]),end(ad[i]),next_f);
}
}
std::vector<int> ans(ad.size());
std::transform(begin(f.back()),end(f.back()),begin(ans),[](char x){
return !x;
});
return ans;
}
