#include "stations.h"
#include <vector>
#include <cassert>
#include <cmath>
#include <algorithm>
namespace std {
template<class Fun>
class y_combinator_result {
Fun fun_;
public:
template<class T>
explicit y_combinator_result(T &&fun): fun_(std::forward<T>(fun)) {}
template<class ...Args>
decltype(auto) operator()(Args &&...args) {
return fun_(std::ref(*this), std::forward<Args>(args)...);
}
};
template<class Fun>
decltype(auto) y_combinator(Fun &&fun) {
return y_combinator_result<std::decay_t<Fun>>(std::forward<Fun>(fun));
}
} // namespace std
std::vector<int> label(int N, int K, std::vector<int> U, std::vector<int> V) {
std::vector<std::vector<int>> adj(N);
for (int e = 0; e < N-1; e++) {
adj[U[e]].push_back(V[e]);
adj[V[e]].push_back(U[e]);
}
std::vector<int> labels(N);
int cur_idx = 0;
std::y_combinator([&](auto self, int cur, int prv, bool dir) -> void {
if (!dir) {
labels[cur] = cur_idx++;
}
for (int nxt : adj[cur]) {
if (nxt == prv) continue;
self(nxt, cur, !dir);
}
if (dir) {
labels[cur] = cur_idx++;
}
})(0, -1, false);
assert(cur_idx-1 <= K);
return labels;
}
int find_next_station(int S, int T, std::vector<int> C) {
assert(S != T);
assert(!C.empty());
if (C[0] < S) {
S = -S;
T = -T;
for (int& c : C) c = -c;
}
assert(C[0] > S);
int par = *std::max_element(C.begin(), C.end());
if (T < S || T >= par) {
return abs(par);
}
int ans = par;
for (int c : C) {
if (c >= T) ans = std::min(ans, c);
}
return abs(ans);
}
