#include "stations.h"
#include <bits/stdc++.h>
using namespace std;
const int N = 1005;
vector<int> tree[N];
int depth[N];
int parent[N][10];
int label_arr[N];
int n;
// Preprocessing for LCA
void dfs(int u, int p) {
parent[u][0] = p;
for (int v : tree[u]) {
if (v != p) {
depth[v] = depth[u] + 1;
dfs(v, u);
}
}
}
void build_lca() {
for (int j = 1; j < 10; ++j) {
for (int i = 0; i < n; ++i) {
if (parent[i][j-1] != -1)
parent[i][j] = parent[parent[i][j-1]][j-1];
else
parent[i][j] = -1;
}
}
}
int get_lca(int u, int v) {
if (depth[u] < depth[v]) swap(u, v);
for (int i = 9; i >= 0; --i) {
if (parent[u][i] != -1 && depth[parent[u][i]] >= depth[v])
u = parent[u][i];
}
if (u == v) return u;
for (int i = 9; i >= 0; --i) {
if (parent[u][i] != -1 && parent[u][i] != parent[v][i]) {
u = parent[u][i];
v = parent[v][i];
}
}
return parent[u][0];
}
int get_dist(int u, int v) {
int lca = get_lca(u, v);
return depth[u] + depth[v] - 2 * depth[lca];
}
// Official functions the judge will call:
vector<int> label(int N, int K, vector<int> U, vector<int> V) {
n = N;
for (int i = 0; i < (int)U.size(); ++i) {
int u = U[i], v = V[i];
tree[u].push_back(v);
tree[v].push_back(u);
}
memset(parent, -1, sizeof(parent));
dfs(0, -1);
build_lca();
// Trivial labels: just label station i as i+1 (to be between 1 and n)
for (int i = 0; i < n; ++i) {
label_arr[i] = i + 1;
}
vector<int> result;
for (int i = 0; i < n; ++i) {
result.push_back(label_arr[i]);
}
return result;
}
int find_next_station(int s, int t, vector<int> neighbors) {
int real_s = s;
int real_t = t;
int best_neighbor = neighbors[0];
int best_dist = get_dist(real_s, real_t);
for (int nb_label : neighbors) {
int real_nb = nb_label - 1; // Because we labeled station i with i+1
int d = get_dist(real_nb, real_t);
if (d < best_dist) {
best_dist = d;
best_neighbor = nb_label;
}
}
return best_neighbor;
}
