#include "triples.h"
long long count_triples(std::vector<int> H) {
std::vector<std::vector<int>> asd (H.size());
for (int i=0; i<H.size(); i++)
asd[H[i]].push_back(i);
std::vector<std::vector<int>> cache(H.size());
for (int i=0; i<H.size(); i++)
if (i-H[i] >= 0)
cache[i-H[i]].push_back(i);
long long ans = 0;
for (int i=0;i<H.size(); i++) {
if (i + H[i] >= H.size()) continue;
int sub_peak_index = H[i] + i;
// One case - special case
for (int x: cache[i+H[i]]) { // cba
if (H[x] + i >= H.size()) continue;
int sub_peak = H[x] + i;
if (H[sub_peak] == H[i] + H[x])
ans++;
}
if (i + H[sub_peak_index] < H.size()) { // abc
int last = i + H[sub_peak_index];
if (H[i] + H[last] == H[sub_peak_index] && H[i] != H[last])
ans++;
}
if (i + H[sub_peak_index] < H.size()) { // cab
int last = i + H[sub_peak_index];
if (H[i] == H[sub_peak_index] + H[last])
ans++;
}
if (sub_peak_index - H[sub_peak_index] >= 0) { // bac
int last = H[i] + i - H[sub_peak_index];
if (H[last] + H[sub_peak_index] == H[i] && H[last] != H[sub_peak_index])
ans ++;
}
for (int sub_peak_index: cache[i]) { // bca
if (sub_peak_index - H[i] < 0) continue;
int last = sub_peak_index - H[i];
if (H[last] + H[i] == H[sub_peak_index])
ans++;
}
if (sub_peak_index + H[sub_peak_index] <= H.size()) { // acb
int last = sub_peak_index + H[sub_peak_index];
if (H[last] == H[i] + H[sub_peak_index] && H[i] != H[sub_peak_index])
ans ++;
}
}
return ans;
}
std::vector<int> construct_range(int M, int K) {
std::vector<int> ans;
ans.push_back(1);
for (int i=1; i<M/2; i++)
ans.push_back(i);
while (ans.size() != M) {
ans.push_back(std::max(1, ans.back() - 1));
}
return ans;
}
