#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);
// Compute side mappings
long long ans = 0;
for (int main_peak=1; main_peak<H.size(); main_peak++) {
for (int j=0; j<asd[main_peak].size(); j++) {
int peak_index = asd[main_peak][j];
// check upper
if (peak_index + main_peak < H.size()) {
int sub_peak_index = peak_index + main_peak;
int sub_peak = H[sub_peak_index];
// check from sub
if (sub_peak_index + sub_peak < H.size())
if (H[sub_peak_index + sub_peak] == main_peak + sub_peak)
ans++;
if (peak_index + sub_peak < H.size())
if (H[peak_index + sub_peak] + sub_peak == main_peak)
ans++;
if (peak_index - sub_peak > peak_index)
if (H[peak_index - sub_peak] == main_peak + sub_peak)
ans++;
if (sub_peak_index - sub_peak > peak_index)
if (H[sub_peak_index - sub_peak] + sub_peak == main_peak)
ans++;
}
}
}
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);
// Compute opposite mappings
for (int i=0;i<H.size(); i++) {
if (i + H[i] >= H.size()) continue;
for (int x: cache[i+H[i]])
if (H[H[x] + i] == H[i] + H[x])
ans++;
}
return ans;
}
std::vector<int> construct_range(int M, int K) {
return {1, 1, 1};
}
