#include "migrations.h"
#include <bits/stdc++.h>
using namespace std;
/*
General solution (works for both subtasks):
During the first run (send_message calls in increasing i):
- We build the tree online and maintain binary-lifting parents (up[k][i]) and depths.
- We maintain the current diameter endpoints (A, B) and its length D.
For each new node i with parent p:
* Compute dist(i, A) and dist(i, B) using LCA.
* If dist(i, A) > D, set (B = i) and update D.
* else if dist(i, B) > D, set (A = i) and update D.
- We send exactly TWO messages total:
* at i == N-2 we send A+1
* at i == N-1 we send B+1
(Each is ≤ N ≤ 10000 ≤ 20000, and M = 2 ≤ 50.)
During the second run (longest_path):
- Read the last two non-zero entries of S (order matters):
* the penultimate non-zero is A+1, the last non-zero is B+1.
- Return (A, B).
This yields the true diameter for any tree (no need for the "root is an endpoint" assumption).
*/
// ---------------- Persistent state for the first run ----------------
namespace {
constexpr int MAXN = 10000 + 5;
constexpr int LOG = 15; // since 2^14 = 16384 > 10000
static bool initialized = false;
static int depth_arr[MAXN];
static int up[LOG][MAXN]; // up[k][v] = 2^k-th ancestor of v, or -1
static int A = 0, B = 0; // current diameter endpoints
static int D = 0; // current diameter length
inline int lca(int u, int v) {
if (u == -1 || v == -1) return u == -1 ? v : u;
if (depth_arr[u] < depth_arr[v]) swap(u, v);
int diff = depth_arr[u] - depth_arr[v];
for (int k = 0; k < LOG; ++k) if (diff & (1 << k)) u = up[k][u];
if (u == v) return u;
for (int k = LOG - 1; k >= 0; --k) {
if (up[k][u] != up[k][v]) {
u = up[k][u];
v = up[k][v];
}
}
return up[0][u];
}
inline int dist(int u, int v) {
int w = lca(u, v);
return depth_arr[u] + depth_arr[v] - 2 * depth_arr[w];
}
inline void add_node(int i, int p) {
// Set depth and ancestors for node i
depth_arr[i] = depth_arr[p] + 1;
up[0][i] = p;
for (int k = 1; k < LOG; ++k) {
int mid = up[k - 1][i];
up[k][i] = (mid == -1 ? -1 : up[k - 1][mid]);
}
// Try to improve diameter with endpoint A or B
int dA = dist(i, A);
if (dA > D) {
B = i;
D = dA;
return;
}
int dB = dist(i, B);
if (dB > D) {
A = i;
D = dB;
return;
}
}
}
// ---------------- Research team procedure (first run) ----------------
int send_message(int N, int i, int Pi) {
if (!initialized) {
initialized = true;
// Initialize root (node 0)
depth_arr[0] = 0;
for (int k = 0; k < LOG; ++k) up[k][0] = -1;
// Start diameter as (0,0)
A = 0; B = 0; D = 0;
}
// Add node i with parent Pi
add_node(i, Pi);
// Send two messages only at the last two sites
if (i == N - 2) {
return A + 1; // encode A
}
if (i == N - 1) {
return B + 1; // encode B
}
return 0;
}
// ---------------- Museum procedure (second run) ----------------
std::pair<int,int> longest_path(std::vector<int> S) {
// Extract the last two non-zero messages: penultimate -> A+1, last -> B+1
int N = (int)S.size();
int last_idx = -1, penult_idx = -1;
for (int i = 1; i < N; ++i) {
if (S[i] != 0) {
penult_idx = last_idx;
last_idx = i;
}
}
if (last_idx == -1) {
// No messages found; fallback to trivial answer
return {0, 0};
}
if (penult_idx == -1) {
// Only one message found; interpret it as B+1 with A assumed 0 as a best-effort fallback
int b = S[last_idx] - 1;
if (b < 0) b = 0;
return {0, b};
}
int a = S[penult_idx] - 1;
int b = S[last_idx] - 1;
if (a < 0) a = 0;
if (b < 0) b = 0;
return {a, b};
}
