#include "stations.h"
//#include "stub.cpp"
#include <bits/stdc++.h>
using namespace std;
#define ll long long
#define pii pair<int,int>
#define all(x) x.begin(),x.end()
#define _ ios::sync_with_stdio(0); cin.tie(0); cout.tie(0);
std::vector<int> label(int n, int k, std::vector<int> u, std::vector<int> v) {
vector<vector<int>> adj(n);
for(int i=0;i<n-1;i++){
adj[u[i]].push_back(v[i]);
adj[v[i]].push_back(u[i]);
}
vector<int> st(n),en(n);
int timer=-1;
function<void(int,int)> dfs=[&](int v,int p){
st[v]=++timer;
for(auto u:adj[v]){
if(u==p) continue;
dfs(u,v);
}
en[v]=timer;
};
dfs(0,-1);
auto get=[&](int a,int b){
int pre=(1000+1000-a+1)*a/2;
return pre+b;
};
vector<int> label(n);
for(int i=0;i<n;i++){
label[i]=get(st[i],en[i]);
//cout<<i<<' '<<label[i]<<'\n';
}
return label;
}
int find_next_station(int s, int t, std::vector<int> C) {
auto get=[&](int x){
int i=0;
while((1000+1000-i+1)*i/2<=x) i++;
i--;
int j=x-(1000+1000-i+1)*i/2;
return make_pair(i,j);
};
auto [a,b]=get(s);
auto [c,d]=get(t);
if(a<=c and d<=b){
//cout<<"b"<<'\n';
for(auto v:C){
auto [e,f]=get(v);
if(e<=c and d<=f and a<=e and f<=b) return v;
}
}
else{
for(auto v:C){
auto [e,f]=get(v);
if(e<=a and b<=f) return v;
}
}
assert(false);
}
