#include "migrations.h"
#include <vector>
#include <utility>
#include <algorithm>
/*
Online diameter tracker (general tree) with robust per-test reinit.
Key fix: reinitialize state at the *start of each test case*, which is the
call with i == 1. The previous version only initialized once per process,
so multiple test cases in the same run could reuse stale state and produce
off-by-one / incorrect diameters.
Strategy:
- Maintain depths and a root-sticky binary-lifting table up[LOG][v].
Root 0 has up[k][0] = 0 for all k.
- Maintain current diameter endpoints (A, B) and its length D online:
* On adding node i with parent p, compute dist(i, A) and dist(i, B).
* If dist(i, A) > D => (B = i, D = dist(i, A));
else if dist(i, B) > D => (A = i, D = dist(i, B)).
- Send exactly two messages:
* at i == N-2 : send A+1
* at i == N-1 : send B+1
(≤ 50 messages, values ≤ N ≤ 10000 ≤ 20000.)
Museum:
- Read the last two non-zero messages: penultimate => A+1, last => B+1.
- Return (A, B).
*/
namespace {
constexpr int MAXN = 10000 + 5;
constexpr int LOG = 15; // 2^14 = 16384 > 10000
// Persistent across calls but re-inited per test case (when i==1).
static bool has_any_init = false;
static int depth_arr[MAXN];
static int up[LOG][MAXN]; // up[k][v] = 2^k-th ancestor, root-sticky
static int A = 0, B = 0; // diameter endpoints
static int D = 0; // diameter length
inline void init_for_test() {
// Root initialization
depth_arr[0] = 0;
for (int k = 0; k < LOG; ++k) up[k][0] = 0;
// Reset diameter
A = 0; B = 0; D = 0;
has_any_init = true;
}
inline int lca(int u, int v) {
if (depth_arr[u] < depth_arr[v]) std::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) {
depth_arr[i] = depth_arr[p] + 1;
up[0][i] = p;
for (int k = 1; k < LOG; ++k)
up[k][i] = up[k - 1][ up[k - 1][i] ];
// Try to extend diameter using new node i
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 (first run): called for i = 1..N-1
int send_message(int N, int i, int Pi) {
// Reinitialize state at the start of each test case (first call has i==1).
if (i == 1 || !has_any_init) {
init_for_test();
}
add_node(i, Pi);
if (i == N - 2) return A + 1; // encode A
if (i == N - 1) return B + 1; // encode B
return 0;
}
// Museum (second run): decode the two endpoints from the last two messages
std::pair<int,int> longest_path(std::vector<int> S) {
int N = (int)S.size();
int last = -1, penult = -1;
for (int i = 1; i < N; ++i) {
if (S[i] != 0) { penult = last; last = i; }
}
if (last == -1) return {0, 0}; // fallback
if (penult == -1) {
int b = S[last] - 1; if (b < 0) b = 0;
return {0, b};
}
int a = S[penult] - 1; if (a < 0) a = 0;
int b = S[last] - 1; if (b < 0) b = 0;
return {a, b};
}
