#include "migrations.h"
#include <vector>
#include <utility>
#include <algorithm>
/*
Robust diameter-in-a-tree strategy (works for general case):
First run (send_message):
- Build the tree online since P[i] < i and i increases.
- Maintain binary lifting table up[LOG][v] with root 0 having up[k][0] = 0
(so we never use -1 and avoid out-of-bounds during LCA).
- Track depths.
- Maintain current diameter endpoints (A, B) and its length D:
* For each new node i, 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 (≤ 50 and values ≤ N ≤ 10000 ≤ 20000):
* at i == N-2: send A+1
* at i == N-1: send B+1
Second run (longest_path):
- Read the last two non-zero S entries: penultimate => A+1, last => B+1.
- Return (A, B).
Fix vs previous attempt:
- We now keep up[k][0] = 0 for all k and set up via up[k][v] = up[k-1][ up[k-1][v] ].
This avoids ever assigning -1 and eliminates potential out-of-bounds in LCA,
which could cause an off-by-one error in the diameter on adversarial trees.
*/
namespace {
constexpr int MAXN = 10000 + 5;
constexpr int LOG = 15; // 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 (root-sticky)
static int A = 0, B = 0; // current diameter endpoints
static int D = 0; // current diameter length
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
depth_arr[i] = depth_arr[p] + 1;
// ancestors (root-sticky)
up[0][i] = p;
for (int k = 1; k < LOG; ++k) {
up[k][i] = up[k - 1][ up[k - 1][i] ];
}
// attempt to extend diameter
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: called N-1 times with i = 1..N-1
int send_message(int N, int i, int Pi) {
if (!initialized) {
initialized = true;
// Root initialization
depth_arr[0] = 0;
for (int k = 0; k < LOG; ++k) up[k][0] = 0;
// Start with trivial diameter
A = 0; B = 0; D = 0;
}
// Add the current node
add_node(i, Pi);
// Send two messages only at the last two indices
if (i == N - 2) return A + 1; // encode A in [1..N]
if (i == N - 1) return B + 1; // encode B in [1..N]
return 0;
}
// Museum: called once in the second run
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) {
// No messages (shouldn't happen with our sender); fallback
return {0, 0};
}
if (penult == -1) {
// Only one message; fallback assuming one endpoint is 0
int b = S[last] - 1;
if (b < 0) b = 0;
return {0, b};
}
int a = S[penult] - 1;
int b = S[last] - 1;
if (a < 0) a = 0;
if (b < 0) b = 0;
return {a, b};
}
