#include "migrations.h"
#include <vector>
#include <algorithm>
#include <cmath>
namespace {
// These variables will persist between calls to send_message
int N_val;
// Data structures for representing the tree and calculating LCA
std::vector<int> parent;
std::vector<int> depth;
std::vector<std::vector<int>> up;
int LOG_N;
// Variables to track the current diameter
int endpoint1 = 0, endpoint2 = 0;
// Variables for the communication protocol
int start_messaging_idx;
bool init_sent = false;
bool second_endpoint_pending = false;
int pending_value = 0;
int first_msg_idx = -1;
}
/**
* @brief Initializes the data structures for LCA calculations.
* This is called only once at the beginning of the simulation.
* @param N The total number of sites.
*/
void init_lca(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; // The root's parent is itself
for(int k = 0; k < LOG_N; ++k) up[0][k] = 0;
}
/**
* @brief Adds a new node to the tree structure for LCA calculations.
* @param i The index of the new node.
* @param p The index of its parent.
*/
void add_node_lca(int i, int p) {
parent[i] = p;
depth[i] = depth[p] + 1;
up[i][0] = p;
for (int k = 1; k < LOG_N; ++k) {
up[i][k] = up[up[i][k-1]][k-1];
}
}
/**
* @brief Finds the Lowest Common Ancestor (LCA) of two nodes using binary lifting.
* @param u First node.
* @param v Second node.
* @return The LCA of u and v.
*/
int get_lca(int u, int v) {
if (depth[u] < depth[v]) std::swap(u, v);
// Bring u to the same depth as v
for (int k = LOG_N - 1; k >= 0; --k) {
if ((depth[u] - (1 << k)) >= depth[v]) {
u = up[u][k];
}
}
if (u == v) return u;
// Move u and v up together until their parents are the same
for (int k = LOG_N - 1; k >= 0; --k) {
if (up[u][k] != up[v][k]) {
u = up[u][k];
v = up[v][k];
}
}
return parent[u];
}
/**
* @brief Calculates the distance between two nodes in the tree.
* @param u First node.
* @param v Second node.
* @return The distance between u and v.
*/
int get_dist(int u, int v) {
if (u == -1 || v == -1) return -1;
int lca = get_lca(u, v);
return depth[u] + depth[v] - 2 * depth[lca];
}
/**
* @brief This function is called for each new site to determine if a message should be sent.
* @param N The total number of sites.
* @param i The index of the current site being added.
* @param Pi The parent of the current site.
* @return The message to be sent (0 if no message).
*/
int send_message(int N, int i, int Pi) {
// Initialization on the first call
if (i == 1) {
init_lca(N);
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;
}
// Add the new node to our tree representation
add_node_lca(i, Pi);
// Calculate the current diameter length
int current_diameter_len = get_dist(endpoint1, endpoint2);
// Calculate distances from the new node to the current diameter endpoints
int dist_i_e1 = get_dist(i, endpoint1);
int dist_i_e2 = get_dist(i, endpoint2);
bool diameter_changed = false;
int other_endpoint = -1;
// Check if the new node extends the diameter
if (dist_i_e1 > current_diameter_len || dist_i_e2 > current_diameter_len) {
diameter_changed = true;
if (dist_i_e1 > dist_i_e2) {
other_endpoint = endpoint1;
endpoint2 = i;
} else {
other_endpoint = endpoint2;
endpoint1 = i;
}
}
// We only send messages for the last 50 nodes to stay within the limit
if (i < start_messaging_idx) {
return 0;
}
int message_to_send = 0;
// Logic for sending the initial diameter information
if (!init_sent) {
if (diameter_changed) {
// If the diameter changes on our first messaging attempt, send the update
message_to_send = other_endpoint + 1;
} else {
// Otherwise, send the first endpoint and queue the second
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 we are waiting to send the second endpoint
if (diameter_changed) {
// An update takes precedence
message_to_send = other_endpoint + 1;
} else {
// Send the queued second endpoint
message_to_send = pending_value;
}
second_endpoint_pending = false;
} else {
// After initialization, only send messages on diameter changes
if (diameter_changed) {
message_to_send = other_endpoint + 1;
}
}
return message_to_send;
}
/**
* @brief This function is called once after all `send_message` calls.
* It reconstructs the diameter from the received messages.
* @param S A vector containing the messages sent. S[i] is the message from node i.
* @return A pair of integers representing the endpoints of the diameter.
*/
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]});
}
}
// No messages, diameter is (0,0)
if (msgs.empty()) {
return {0, 0};
}
// Only one message, it's an update.
if (msgs.size() == 1) {
return {msgs[0].first, msgs[0].second - 1};
}
// Two messages, could be initialization or an update.
if (msgs.size() == 2) {
int i1 = msgs[0].first;
int v1 = msgs[0].second;
int i2 = msgs[1].first;
int v2 = msgs[1].second;
// Heuristic to differentiate: an initialization value for the second endpoint (v2-1)
// cannot be equal to the index of the first message (i1), as it must be a node
// that existed before i1.
if (v2 == v1 || v2 - 1 == i1) {
return {i2, v2 - 1};
} else {
return {v1 - 1, v2 - 1};
}
}
// For more than two messages, the last one represents the final diameter.
return {msgs.back().first, msgs.back().second - 1};
}
