#include <vector>
#include <utility>
#include <algorithm>
// Global state variables
static int _N = 0;
static std::vector<int> _depth;
static std::vector<std::vector<int>> _up;
static int _LOG = 0;
static int _a = 0;
static int _b = 0;
static int _diam = 0;
static int _a2 = -1; // using -1 for Python's None
static int _b2 = -1;
static int _b_at_N3 = -1;
void _reinit(int N) {
_N = N;
// Equivalent to Python's N.bit_length() + 1
int temp = N;
int bits = 0;
while (temp > 0) {
bits++;
temp >>= 1;
}
_LOG = bits + 1;
_depth.assign(N, 0);
_up.assign(_LOG, std::vector<int>(N, 0));
_a = 0;
_b = 0;
_diam = 0;
_a2 = -1;
_b2 = -1;
_b_at_N3 = -1;
}
int _lca(int u, int v) {
if (_depth[u] < _depth[v]) {
std::swap(u, v);
}
int diff = _depth[u] - _depth[v];
int bit = 0;
while (diff > 0) {
if (diff & 1) {
u = _up[bit][u];
}
diff >>= 1;
bit++;
}
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];
}
int _dist(int u, int v) {
int w = _lca(u, v);
return _depth[u] + _depth[v] - 2 * _depth[w];
}
int send_message(int N, int i, int Pi) {
// Initialise a fresh test case at the first call (i == 1)
if (i == 1) {
_reinit(N);
}
// ---- insert new leaf i ----
int d = _depth[Pi] + 1;
_depth[i] = d;
_up[0][i] = Pi;
for (int k = 1; k < _LOG; ++k) {
int anc = _up[k - 1][i];
_up[k][i] = _up[k - 1][anc];
}
// ---- maintain the current diameter (a,b) ----
int da = _dist(i, _a);
int db = _dist(i, _b);
if (da > _diam) {
_b = i;
_diam = da;
} else if (db > _diam) {
_a = i;
_diam = db;
}
// ---- decide which messages to send (only three are needed) ----
if (_N >= 4 && i == _N - 3 && i >= 1) {
_b_at_N3 = _b;
return _b + 1;
} else if (i == _N - 2 && i >= 1) {
_a2 = _a;
_b2 = _b;
return _a + 1;
} else if (i == _N - 1 && i >= 1) {
int a_final = _a;
int b_final = _b;
int a2 = _a2;
int b2 = _b2;
int b0 = _b_at_N3;
int flag = 0;
// Determine which case we are in and assign the flag.
if (a_final == a2 && b_final == b2) {
if (b2 == _N - 2) {
flag = 2;
} else {
flag = 1;
}
} else if (a_final == a2 && b_final == _N - 1) {
flag = 3;
} else if (a_final == _N - 1 && b_final == b2) {
if (b2 == _N - 2) {
flag = 5;
} else {
flag = 4;
}
} else {
flag = 1;
}
return flag;
} else {
return 0;
}
}
std::pair<int, int> longest_path(std::vector<int> S) {
int N = S.size();
if (N <= 1) {
return {0, 0};
}
if (N == 2) {
return {0, 1};
}
if (N == 3) {
return {1, 2};
}
int a_msg = S[N - 2];
int b_msg = S[N - 3];
int flag = S[N - 1];
int a = a_msg - 1;
int b0 = b_msg - 1;
int u = 0, v = 0;
if (flag == 1) {
u = a;
v = b0;
} else if (flag == 2) {
u = a;
v = N - 2;
} else if (flag == 3) {
u = a;
v = N - 1;
} else if (flag == 4) {
u = b0;
v = N - 1;
} else if (flag == 5) {
u = N - 2;
v = N - 1;
} else {
u = a;
v = b0;
}
return {u, v};
}
