#include "migrations.h"
#include <vector>
#include <algorithm>
#include <cmath>
namespace {
// These variables will persist across calls to send_message
int N_val;
// For LCA and distance calculation
std::vector<int> parent;
std::vector<int> depth;
std::vector<std::vector<int>> up;
int LOG_N;
// For diameter tracking
int endpoint1 = 0, endpoint2 = 0;
// For communication protocol
int start_messaging_idx;
bool init_sent = false;
bool second_endpoint_pending = false;
int pending_value = 0;
int first_msg_idx = -1;
}
// Helper function to calculate distance between two nodes using LCA
int get_dist(int u, int v) {
if (depth[u] < depth[v]) std::swap(u, v);
for (int k = LOG_N - 1; k >= 0; --k) {
if (depth[u] - (1 << k) >= depth[v]) {
u = up[u][k];
}
}
if (u == v) {
// This is a direct ancestor-descendant path, distance is difference in depth.
// Re-calculating depth difference as u or v might have been swapped.
return (depth[u] > depth[v] ? depth[u] : depth[v]) - (depth[u] < depth[v] ? depth[u] : depth[v]);
}
for (int k = LOG_N - 1; k >= 0; --k) {
if (up[u][k] != up[v][k]) {
u = up[u][k];
v = up[v][k];
}
}
int lca = parent[u];
return depth[u] + depth[v] - 2 * depth[lca];
}
void initialize(int N) {
N_val = N;
parent.assign(N, -1);
depth.assign(N, 0);
LOG_N = 0;
while ((1 << LOG_N) <= N) {
LOG_N++;
}
up.assign(N, std::vector<int>(LOG_N, 0));
parent[0] = 0;
for(int k=0; k<LOG_N; ++k) up[0][k] = 0;
endpoint1 = 0;
endpoint2 = 0;
start_messaging_idx = std::max(1, N - 50);
init_sent = false;
second_endpoint_pending = false;
pending_value = 0;
first_msg_idx = -1;
}
int send_message(int N, int i, int Pi) {
if (i == 1) {
initialize(N);
}
// Update tree structure for LCA
parent[i] = Pi;
depth[i] = depth[Pi] + 1;
up[i][0] = Pi;
for (int k = 1; k < LOG_N; ++k) {
up[i][k] = up[up[i][k-1]][k-1];
}
// Update diameter endpoints
int dist_e1_e2 = get_dist(endpoint1, endpoint2);
int dist_i_e1 = get_dist(i, endpoint1);
int dist_i_e2 = get_dist(i, endpoint2);
bool diameter_changed = false;
int other_endpoint = -1;
if (dist_i_e1 > dist_e1_e2 && dist_i_e1 >= dist_i_e2) {
other_endpoint = endpoint1;
endpoint2 = i;
diameter_changed = true;
} else if (dist_i_e2 > dist_e1_e2 && dist_i_e2 > dist_i_e1) {
other_endpoint = endpoint2;
endpoint1 = i;
diameter_changed = true;
}
if (i < start_messaging_idx) {
return 0;
}
int message_to_send = 0;
if (!init_sent) {
if (diameter_changed) {
message_to_send = other_endpoint + 1;
} else {
message_to_send = endpoint1 + 1;
first_msg_idx = i;
second_endpoint_pending = true;
pending_value = endpoint2 + 1;
}
init_sent = true;
} else if (second_endpoint_pending) {
if (diameter_changed) {
message_to_send = other_endpoint + 1;
} else {
message_to_send = pending_value;
}
second_endpoint_pending = false;
} else {
if (diameter_changed) {
message_to_send = other_endpoint + 1;
}
}
return message_to_send;
}
std::pair<int, int> longest_path(std::vector<int> S) {
std::vector<std::pair<int, int>> msgs;
for (size_t i = 0; i < S.size(); ++i) {
if (S[i] != 0) {
msgs.push_back({(int)i, S[i]});
}
}
if (msgs.empty()) {
return {0, 0};
}
if (msgs.size() == 1) {
return {msgs[0].first, msgs[0].second - 1};
}
if (msgs.size() == 2) {
int i1 = msgs[0].first;
int v1 = msgs[0].second;
int i2 = msgs[1].first;
int v2 = msgs[1].second;
if (v2 == v1 || v2 - 1 == i1) {
return {i2, v2 - 1};
} else {
return {v1 - 1, v2 - 1};
}
}
return {msgs.back().first, msgs.back().second - 1};
}
