#include "migrations.h"
#include <bits/stdc++.h>
using namespace std;
static bool initialized = false;
static int N_global = 0;
static const int LOGV = 15; // 2^14 = 16384 > 10000
static vector<int> parentArr;
static vector<int> depthArr;
static vector<array<int, LOGV>> upArr;
static int endA = 0, endB = 0, curDist = 0; // current diameter endpoints and length
static int repl_Nminus2 = -1; // replacement code for node N-2 (0/1/2)
static int lca(int u, int v) {
if (depthArr[u] < depthArr[v]) swap(u, v);
int diff = depthArr[u] - depthArr[v];
for (int k = 0; diff; ++k) {
if (diff & 1) u = upArr[u][k];
diff >>= 1;
}
if (u == v) return u;
for (int k = LOGV - 1; k >= 0; --k) {
if (upArr[u][k] != upArr[v][k]) {
u = upArr[u][k];
v = upArr[v][k];
}
}
return upArr[u][0];
}
static int dist(int u, int v) {
int w = lca(u, v);
return depthArr[u] + depthArr[v] - 2 * depthArr[w];
}
/* return code of what happened when node x is added:
0 – no change,
1 – endpoint A is replaced by x,
2 – endpoint B is replaced by x */
static int update_diameter(int x) {
int da = dist(x, endA);
int db = dist(x, endB);
if (da > curDist) {
endB = x;
curDist = da;
return 2;
} else if (db > curDist) {
endA = x;
curDist = db;
return 1;
}
return 0;
}
int send_message(int N, int i, int Pi) {
if (!initialized) {
N_global = N;
parentArr.assign(N, -1);
depthArr.assign(N, 0);
upArr.assign(N, array<int, LOGV>{});
for (int k = 0; k < LOGV; ++k) upArr[0][k] = 0; // root
endA = endB = 0;
curDist = 0;
repl_Nminus2 = -1;
initialized = true;
}
// add node i
parentArr[i] = Pi;
depthArr[i] = depthArr[Pi] + 1;
upArr[i][0] = Pi;
for (int k = 1; k < LOGV; ++k)
upArr[i][k] = upArr[ upArr[i][k - 1] ][k - 1];
// handle tiny N specially
if (N_global == 2) {
update_diameter(i);
return 0;
}
if (N_global == 3) {
if (i == 1) {
update_diameter(i);
return 0;
} else { // i == 2, last node
update_diameter(i);
int a = endA, b = endB;
if (a > b) swap(a, b);
int code = a * N_global + b + 1; // fits into [1,9] for N=3
return code;
}
}
// general case N >= 4
if (i == N_global - 3) {
// after processing this node, send endpoint A
update_diameter(i);
return endA + 1;
}
if (i == N_global - 2) {
// store current endpoint B before possible change
int storedB = endB;
repl_Nminus2 = update_diameter(i);
return storedB + 1;
}
if (i == N_global - 1) {
int repl_last = update_diameter(i);
int code = repl_Nminus2 * 3 + repl_last; // 0..8
return code + 1; // 1..9
}
// all other nodes
update_diameter(i);
return 0;
}
std::pair<int, int> longest_path(std::vector<int> S) {
int N = (int)S.size();
if (N == 2) return {0, 1};
if (N == 3) {
int code = -1;
for (int i = 0; i < N; ++i) if (S[i] != 0) { code = S[i] - 1; break; }
int a = code / N;
int b = code % N;
return {a, b};
}
// N >= 4
int a = S[N - 3] - 1;
int b = S[N - 2] - 1;
int code = S[N - 1] - 1;
int r2 = code / 3;
int r3 = code % 3;
if (r2 == 1) a = N - 2;
else if (r2 == 2) b = N - 2;
if (r3 == 1) a = N - 1;
else if (r3 == 2) b = N - 1;
return {a, b};
}
