#include "migrations.h"
#include <vector>
#include <queue>
#include <utility>
#include <algorithm>
using namespace std;
/*-------------------------------------------------------------
Global data kept by the research team (persist between calls)
-------------------------------------------------------------*/
static int GLOBAL_N = -1; // size of the current test case
static vector<vector<int>> adj; // adjacency list
static vector<int> parent_; // parent of each vertex (P[i])
static vector<int> depth_; // depth from root 0
static bool prepared = false; // have we already computed the answer?
static int answer_a = -1; // first endpoint of diameter
static int answer_b = -1; // second endpoint
/*-------------------------------------------------------------
Breadth‑first search – returns (farthest_vertex, distance[])
-------------------------------------------------------------*/
static pair<int, vector<int>> bfs_farthest(int start)
{
vector<int> dist(GLOBAL_N, -1);
queue<int> q;
q.push(start);
dist[start] = 0;
int far = start;
while (!q.empty()) {
int v = q.front(); q.pop();
for (int nb : adj[v]) {
if (dist[nb] == -1) {
dist[nb] = dist[v] + 1;
q.push(nb);
if (dist[nb] > dist[far]) far = nb;
}
}
}
return {far, std::move(dist)};
}
/*-------------------------------------------------------------
Called by the research team for every site i (1 … N‑1)
-------------------------------------------------------------*/
int send_message(int N, int i, int Pi)
{
// first call – allocate structures
if (GLOBAL_N != N) {
GLOBAL_N = N;
adj.assign(N, {});
parent_.assign(N, -1);
depth_.assign(N, 0);
prepared = false;
}
// add the new edge i – Pi
adj[i].push_back(Pi);
adj[Pi].push_back(i);
parent_[i] = Pi;
depth_[i] = depth_[Pi] + 1;
// When the whole tree becomes known (i == N-2) we can compute the diameter
if (i == N - 2 && !prepared) {
// first BFS from any vertex, e.g. 0
auto first = bfs_farthest(0);
int a = first.first;
// second BFS from a
auto second = bfs_farthest(a);
int b = second.first;
answer_a = a; // store endpoints
answer_b = b;
prepared = true; // we will use them for the last two messages
}
// send the two messages at the last two sites
if (i == N - 2 && prepared) {
// send a+1 (must be between 1 and 20000)
return answer_a + 1;
}
if (i == N - 1 && prepared) {
return answer_b + 1;
}
// no message otherwise
return 0;
}
/*-------------------------------------------------------------
Called once after the whole run – the museum receives S[]
-------------------------------------------------------------*/
std::pair<int,int> longest_path(std::vector<int> S)
{
int N = (int)S.size(); // S[0] … S[N‑1]
int a = -1, b = -1;
for (int i = 0; i < N; ++i) {
if (S[i] != 0) {
if (a == -1) a = S[i] - 1;
else b = S[i] - 1;
}
}
// safety: if for some reason a or b was not set, fall back to (0,0)
if (a == -1) a = 0;
if (b == -1) b = 0;
return {a, b};
}
