#include "stations.h"
#include <bits/stdc++.h>
using namespace std;
typedef vector<int> vi;
typedef vector<vi> vvi;
#define ALL(v) v.begin(), v.end()
vvi el;
bitset<1001> visited{};
vi counts{};
vi labelsOut{};
int countStations(int n) {
visited[n] = true;
int count = 1;
for (int to : el[n]) {
if (!visited[to]) {
count += countStations(to);
}
}
counts[n] = count;
return count;
}
void assign(int n, int u, int l, bool upperbound) {
int c = 0;
for (int to : el[n]) {
if (labelsOut[to] == -1) {
labelsOut[to] = u - c - (upperbound ? 0 : counts[to] - 1);
assign(to, u - c - (upperbound ? 1 : 0), u - c - counts[to] + 1 + (upperbound ? 0 : 1), !upperbound);
c += counts[to];
}
}
}
std::vector<int> label(int n, int k, std::vector<int> u, std::vector<int> v) {
el.resize(n);
counts.resize(n);
labelsOut.assign(n, -1);
for (int i = 0; i < n-1; i++) {
el[u[i]].push_back(v[i]);
el[v[i]].push_back(u[i]);
}
countStations(0);
labelsOut[0] = 0;
assign(0, n-1, 1, true);
return labelsOut;
}
int find_next_station(int s, int t, std::vector<int> c) {
bool upperbound = s < c[0];
int cz = c.size();
int prev = (upperbound ? c.back() : c[0]);
if (upperbound) {
if (t < s || t > c[cz-2]) return prev;
auto res = lower_bound(ALL(c), t);
return *res;
} else {
if (t > s || t < c[1]) return prev;
auto res = upper_bound(ALL(c), t);
res--;
return *res;
}
}
