#include <vector>
#include <unordered_map>
#include <utility>
#include <algorithm>
// -------------------------------------------------------------
// Global state kept between successive calls of send_message
// -------------------------------------------------------------
std::vector<int> _depth;
int _max_depth = -1;
int _deepest = -1;
std::unordered_map<int, int> _pending;
bool _encoded = false;
bool _fallback = false;
int _L = 0;
int _suffix_start = 0;
bool _initialized = false;
int send_message(int N, int i, int Pi) {
// ----- initialise on the first call -----
if (!_initialized) {
_depth.assign(N, 0); // depth[0] = 0
_max_depth = -1;
_deepest = -1;
_pending.clear();
_encoded = false;
_fallback = false;
// smallest L with 3**L >= N
_L = 0;
long long power_of_3 = 1;
while (power_of_3 < N) {
_L++;
power_of_3 *= 3;
}
_suffix_start = N - _L;
_initialized = true;
}
// ----- compute depth of the current vertex -----
int d = _depth[Pi] + 1;
_depth[i] = d;
// ----- update deepest vertex seen so far -----
if (d > _max_depth) {
_max_depth = d;
_deepest = i;
}
// ----- already in fallback mode ? -----
if (_fallback) {
// send a message only from the current deepest vertex
return (i == _deepest) ? 1 : 0;
}
// ----- a deeper vertex appears inside the suffix -> start fallback -----
if (_deepest >= _suffix_start && i == _deepest) {
_fallback = true;
_pending.clear(); // cancel everything still pending
return 1; // message for the deepest vertex
}
// ----- first position of the suffix, still normal mode -----
if (i == _suffix_start && !_encoded) {
if (_deepest < _suffix_start) { // deepest is before the suffix -> encode it
_encoded = true;
// base-3 representation, exactly _L digits (most significant first)
std::vector<int> digits;
int x = _deepest;
for (int j = 0; j < _L; ++j) {
digits.push_back(x % 3);
x /= 3;
}
std::reverse(digits.begin(), digits.end());
// first digit is sent now
int first_val = digits[0] + 2; // 2, 3 or 4
// schedule the remaining digits
for (int offset = 1; offset < _L; ++offset) {
_pending[_suffix_start + offset] = digits[offset] + 2;
}
return first_val;
} else {
// deepest already inside the suffix - nothing to encode
return 0;
}
}
// ----- inside the suffix while we are encoding -----
if (_encoded) {
auto it = _pending.find(i);
if (it != _pending.end()) {
int val = it->second;
_pending.erase(it); // scheduled digit (mimicking dict.pop)
return val;
} else {
return 0; // all digits already sent
}
}
// ----- no message to send -----
return 0;
}
// ---------- second run ------------------------------------------------
std::pair<int, int> longest_path(std::vector<int> S) {
int N = S.size();
// recompute L (same as in the first run)
int L = 0;
long long power_of_3 = 1;
while (power_of_3 < N) {
L++;
power_of_3 *= 3;
}
// locate the last non-zero entry
int i_last = -1;
for (int i = N - 1; i >= 0; --i) {
if (S[i] != 0) {
i_last = i;
break;
}
}
// according to the construction a non-zero entry always exists
if (i_last == -1) {
return {0, 0};
}
// fallback case - the special value 1 marks the deepest vertex
if (S[i_last] == 1) {
return {0, i_last};
}
// encoded suffix: collect the whole consecutive suffix of values >= 2
std::vector<int> digits;
int i = i_last;
while (i >= 0 && S[i] >= 2) {
digits.push_back(S[i] - 2); // back to a base-3 digit (0, 1, 2)
i--;
}
std::reverse(digits.begin(), digits.end()); // most-significant first
// rebuild the vertex number
int idx = 0;
for (int d : digits) {
idx = idx * 3 + d;
}
return {0, idx};
}
