#include "triples.h"
#include <bits/stdc++.h>
using namespace std;
static inline uint64_t encode(int i, int j, int k) {
return ( (uint64_t)i << 40 ) | ( (uint64_t)j << 20 ) | (uint64_t)k;
}
void addTriple(int i, int j, int k, unordered_set<uint64_t> &S) {
if (i < j && j < k) S.insert(encode(i, j, k));
}
/* main routine required by the statement */
long long count_triples(std::vector<int> H) {
const int N = (int)H.size();
vector<int> forward(N, -1), backward(N, -1);
vector<vector<int>> forward_to(N), backward_to(N);
for (int i = 0; i < N; ++i) {
int f = i + H[i];
if (f < N) {
forward[i] = f;
forward_to[f].push_back(i);
}
int b = i - H[i];
if (b >= 0) {
backward[i] = b;
backward_to[b].push_back(i);
}
}
unordered_set<uint64_t> triples;
triples.reserve(N * 4);
/* case 1 : a , b , a+b (two forward edges) */
for (int i = 0; i < N; ++i) {
int j = forward[i];
if (j == -1) continue;
int k = forward[j];
if (k == -1) continue;
if ((long long)H[i] + H[j] == H[k]) addTriple(i, j, k, triples);
}
/* case 2 : a , a+b , b (i -> j, middle largest) */
for (int i = 0; i < N; ++i) {
int j = forward[i];
if (j == -1) continue;
int D = H[j];
if (D <= H[i]) continue; // D must be the largest
int k = i + D;
if (k >= N) continue;
if (H[k] == D - H[i]) addTriple(i, j, k, triples);
}
/* case 3 : b , a+b , a (k -> j, middle largest) */
for (int k = 0; k < N; ++k) {
int j = backward[k];
if (j == -1) continue;
int D = H[j];
if (D <= H[k]) continue;
int i = j - (D - H[k]); // i = j - b
if (i < 0) continue;
if (H[i] == D - H[k]) addTriple(i, j, k, triples);
}
/* case 4 : b , a , a+b (j - H[j] = i, then k = j + H[i]) */
for (int j = 0; j < N; ++j) {
int i = backward[j];
if (i == -1) continue;
int hi = H[i];
int k = j + hi;
if (k >= N) continue;
if (H[k] == hi + H[j]) addTriple(i, j, k, triples);
}
/* case 5 : a+b , a , b (i -> k, j = k - H[k]) */
for (int i = 0; i < N; ++i) {
int k = forward[i];
if (k == -1) continue;
int j = backward[k];
if (j <= i || j >= k) continue;
if (backward[j] == i) addTriple(i, j, k, triples);
}
/* case 6 : a+b , b , a (two forward edges to the same k) */
for (int k = 0; k < N; ++k) {
int diff = H[k];
const vector<int> &vec = forward_to[k];
if (vec.empty()) continue;
// store indices in an unordered_set for O(1) existence test
unordered_set<int> setVec(vec.begin(), vec.end());
for (int i : vec) {
int j = i + diff;
if (j < k && setVec.find(j) != setVec.end()) {
// diff = H[k] is already the distance j-i needed
addTriple(i, j, k, triples);
}
}
}
/* pivot case – permutation 4 (b , a+b , a) */
for (int p = 0; p < N; ++p) {
const vector<int> &I = forward_to[p];
const vector<int> &K = backward_to[p];
if (I.empty() || K.empty()) continue;
for (int i : I) {
for (int k : K) {
if (i >= k) continue;
long j = (long)i + k - p;
if (j <= i || j >= k || j >= N) continue;
if (H[(int)j] == k - i) addTriple(i, (int)j, k, triples);
}
}
}
return (long long)triples.size();
}
#include <vector>
#include <algorithm>
using namespace std;
std::vector<int> construct_range(int M, int K) {
vector<int> H;
H.reserve(M);
long long total = 0; // number of mythical triples found so far
while ((int)H.size() < M && total < K) {
int idx = (int)H.size(); // index of the new peak (0‑based)
// The first two positions cannot form a triple yet
if (idx < 2) {
H.push_back(1); // any allowed height
continue;
}
// cnt[v] = how many new triples would appear if we set H[idx] = v
vector<int> cnt(idx + 1, 0); // possible heights are 1..idx
// enumerate all pairs (i, j) with i < j < idx
for (int i = 0; i < idx - 1; ++i) {
int hi = H[i];
for (int j = i + 1; j < idx; ++j) {
int hj = H[j];
int a = j - i; // distance i -> j
int b = idx - j; // distance j -> idx
int c = a + b; // distance i -> idx
int v = -1;
if (hi == a && hj == b) v = c;
else if (hi == a && hj == c) v = b;
else if (hi == b && hj == a) v = c;
else if (hi == b && hj == c) v = a;
else if (hi == c && hj == a) v = b;
else if (hi == c && hj == b) v = a;
if (v >= 1 && v <= idx) {
++cnt[v];
}
}
}
// choose the height giving the most new triples
int best_v = 1;
int best_gain = cnt[1];
for (int v = 2; v <= idx; ++v) {
if (cnt[v] > best_gain) {
best_gain = cnt[v];
best_v = v;
}
}
// if no gain is possible, just pick any legal height (1)
if (best_gain <= 0) {
best_v = 1;
best_gain = 0;
}
H.push_back(best_v);
total += best_gain;
}
// If we stopped because we reached M but still have fewer than K triples,
// we simply return the constructed range (partial credit). In practice,
// with the greedy construction the condition total >= K is reached well
// before M = 5000.
return H;
}
