#include <bits/stdc++.h>
using namespace std;
static std::mt19937 rng(
(unsigned)chrono::steady_clock::now().time_since_epoch().count()
);
static long long count_triples_with_limit(const std::vector<int>& H, long long limit) {
int n = (int)H.size();
long long ans = 0;
// leftmost (span at i)
for (int i = 0; i < n; i++) {
int s = H[i];
int k = i + s;
if (k >= n) continue;
int v = H[k];
if (v >= s) continue;
int target = s - v;
int j1 = i + v;
if (i < j1 && j1 < k && H[j1] == target) {
ans++;
if (limit >= 0 && ans >= limit) return ans;
}
int j2 = k - v;
if (j2 != j1 && i < j2 && j2 < k && H[j2] == target) {
ans++;
if (limit >= 0 && ans >= limit) return ans;
}
}
// rightmost (span at k)
for (int k = 0; k < n; k++) {
int s = H[k];
int i = k - s;
if (i < 0) continue;
int v = H[i];
if (v >= s) continue;
int target = s - v;
int j1 = i + v;
if (i < j1 && j1 < k && H[j1] == target) {
ans++;
if (limit >= 0 && ans >= limit) return ans;
}
int j2 = k - v;
if (j2 != j1 && i < j2 && j2 < k && H[j2] == target) {
ans++;
if (limit >= 0 && ans >= limit) return ans;
}
}
// helper structures for the two middle cases
vector<vector<int>> plus(n), minus(n); // plus[r]: i with i + H[i] = r ; minus[r]: k with k - H[k] = r
for (int idx = 0; idx < n; idx++) {
int h = H[idx];
int rp = idx + h;
if (rp < n) plus[rp].push_back(idx);
int rm = idx - h;
if (rm >= 0) minus[rm].push_back(idx);
}
// middle case A: i + H[i] = j , k - H[k] = j
for (int j = 0; j < n; j++) {
const auto& left_i = plus[j];
const auto& right_k = minus[j];
if (left_i.empty() || right_k.empty()) continue;
unordered_map<int, int> freq;
freq.reserve(left_i.size() * 2 + 1);
for (int i : left_i) {
freq[H[i]]++;
}
int hj = H[j];
for (int k : right_k) {
int hk = H[k];
int need = hj - hk;
if (need > 0) {
auto it = freq.find(need);
if (it != freq.end()) {
ans += it->second;
if (limit >= 0 && ans >= limit) return ans;
}
}
}
}
// middle case B: i + H[k] = j , k - H[i] = j
for (int i = 0; i < n; i++) {
int base = i + H[i];
if (base >= n) continue;
for (int k : minus[base]) {
int hk = H[k];
int j = i + hk;
if (j >= n) continue;
if (H[j] != H[i] + hk) continue;
if (hk == H[i]) continue;
if (k != j + H[i]) continue;
ans++;
if (limit >= 0 && ans >= limit) return ans;
}
}
return ans;
}
long long count_triples(std::vector<int> H) {
return count_triples_with_limit(H, -1);
}
static std::vector<int> greedy_construct(int N, int start0 = -1, int start1 = -1) {
vector<int> H(N, 0);
if (N >= 1) H[0] = (start0 != -1 ? start0 : 1);
if (N >= 2) H[1] = (start1 != -1 ? start1 : 1);
for (int pos = 2; pos < N; pos++) {
vector<int> cnt(N + 1, 0); // possible distances are < N
int best_h = 1;
int best_gain = 0;
for (int i = 0; i < pos - 1; i++) {
int hi = H[i];
for (int j = i + 1; j < pos; j++) {
int hj = H[j];
int a = j - i;
int b = pos - j;
int c = pos - i;
auto upd = [&](int hk) {
if (1 <= hk && hk <= N - 1) {
cnt[hk]++;
if (cnt[hk] > best_gain) {
best_gain = cnt[hk];
best_h = hk;
}
}
};
// six possible matchings
if (hi == a && hj == b) {
upd(c);
continue;
}
if (hi == a && hj == c) {
upd(b);
continue;
}
if (hi == b && hj == a) {
upd(c);
continue;
}
if (hi == b && hj == c) {
upd(a);
continue;
}
if (hi == c && hj == a) {
upd(b);
continue;
}
if (hi == c && hj == b) {
upd(a);
continue;
}
}
}
H[pos] = (best_gain > 0 ? best_h : 1);
}
// clamp to allowed interval
for (int i = 0; i < N; i++) {
if (H[i] < 1) H[i] = 1;
else if (H[i] > N - 1) H[i] = N - 1;
}
return H;
}
static std::vector<int> make_pattern(int N, int a) {
vector<int> H(N, 0);
for (int i = 0; i < N; i++) {
if (i % 2 == 0) {
int pos = i / 2;
H[i] = (pos % 3 == 0 ? 2 * a : a);
} else {
int pos = (i - 1) / 2;
H[i] = (pos % 3 == 0 ? 2 * a : a);
}
}
return H;
}
std::vector<int> construct_range(int M, int K) {
/*
Construct a mountain range with at most M peaks that contains at least K mythical triples.
Converted from the given Python version.
*/
int N = max(3, M);
vector<vector<int>> seeds;
// a = b = 1 patterns (three rotations)
{
vector<int> base;
for (int i = 0; i < (N / 3) + 2; i++) {
base.push_back(1);
base.push_back(1);
base.push_back(2);
}
seeds.emplace_back(base.begin(), base.begin() + N);
}
{
vector<int> base;
for (int i = 0; i < (N / 3) + 2; i++) {
base.push_back(1);
base.push_back(2);
base.push_back(1);
}
seeds.emplace_back(base.begin(), base.begin() + N);
}
{
vector<int> base;
for (int i = 0; i < (N / 3) + 2; i++) {
base.push_back(2);
base.push_back(1);
base.push_back(1);
}
seeds.emplace_back(base.begin(), base.begin() + N);
}
// a = b = 2 pattern (2-2-4)
seeds.push_back(make_pattern(N, 2));
// a = b = 3 pattern (3-3-6) if possible
if (N - 1 >= 6) seeds.push_back(make_pattern(N, 3));
// a = b = 4 pattern (4-4-8)
if (N - 1 >= 8) seeds.push_back(make_pattern(N, 4));
uniform_int_distribution<int> dist_val(1, max(1, N - 1));
uniform_int_distribution<int> dist_idx(0, N - 1);
uniform_real_distribution<double> dist_prob(0.0, 1.0);
// random seeds
for (int t = 0; t < 20; t++) {
vector<int> h(N);
for (int i = 0; i < N; i++) h[i] = dist_val(rng);
seeds.push_back(std::move(h));
}
// greedy seeds with random first two values
for (int t = 0; t < 8; t++) {
int s0 = dist_val(rng);
int s1 = dist_val(rng);
seeds.push_back(greedy_construct(N, s0, s1));
}
// evaluate seeds, keep the best
vector<int> best_H;
long long best_T = -1;
for (const auto& h : seeds) {
long long t = count_triples_with_limit(h, K);
if (t > best_T) {
best_T = t;
best_H = h;
}
if (best_T >= K) break;
}
// hill climbing on the best seed
if (best_T < K) {
vector<int> cur_H = best_H;
long long cur_T = best_T;
int max_steps = 5000;
for (int step = 0; step < max_steps; step++) {
int i = dist_idx(rng);
int old = cur_H[i];
int new_val;
if (dist_prob(rng) < 0.5) {
new_val = dist_val(rng);
} else {
int j = dist_idx(rng);
new_val = abs(i - j);
if (new_val == 0) new_val = 1;
if (new_val >= N) new_val = N - 1;
}
if (new_val == old) continue;
cur_H[i] = new_val;
long long new_T = count_triples_with_limit(cur_H, cur_T + 1);
if (new_T > cur_T) {
cur_T = new_T;
if (new_T > best_T) {
best_T = new_T;
best_H = cur_H;
if (best_T >= K) break;
}
} else {
cur_H[i] = old;
}
if (best_T >= K) break;
}
}
// final sanity: ensure values are in range
vector<int> H_out = best_H;
if ((int)H_out.size() > N) H_out.resize(N);
for (int& val : H_out) {
if (val < 1) val = 1;
else if (val > N - 1) val = N - 1;
}
return H_out;
}
