#include "triples.h"
#include <map>
#include <set>
#include <algorithm>
#include <cmath>
using namespace std;
long long count_triples(vector<int> H) {
vector<int> X, Y;
for (int i=0; i < H.size(); i++) {
X.push_back(i+H[i]);
Y.push_back(i-H[i]);
}
set<int> supX = set<int>(X.begin(), X.end());
map<int, vector<int>> invX;
for (int i=0; i < H.size(); i++) invX[X[i]].push_back(i);
set<tuple<int, int, int>> triples;
for (int i=0; i < H.size(); i++) {
int x = X[i], y = Y[i];
if (x < H.size()) {
int yx = Y[x];
int xx = X[x];
if (i < yx) {
int yyx = Y[yx];
if (yyx == i) triples.insert({i, yx, x});
}
if (xx < H.size()) {
int yxx = Y[xx];
if (yxx == i) triples.insert({i, x, xx});
}
int hx = H[x];
if (i + hx < H.size()) {
if (x == X[i + hx]) triples.insert({i, i + hx, x});
if (x == Y[i + hx]) triples.insert({i, x, i + hx});
}
}
if (y >= 0) {
int hy = H[y];
if (i + hy < H.size()) {
if (y == Y[i + hy]) triples.insert({y, i, i + hy});
}
}
}
for (int x : supX) {
auto xIdx = invX[x];
if (xIdx.size() < sqrt(H.size())) {
for (int j: xIdx)
for (int k: xIdx)
if (j < k) {
int yj = Y[j], yk = Y[k];
if ((yk + yj) % 2 == 0) {
int i = (yj + yk) / 2;
if (i >= 0 && i < H.size())
if (H[i] == (yk - yj) / 2) triples.insert({i, j, k});
}
}
} else {
for (int i=0; i < H.size(); i++) {
if ((x + X[i]) % 2 == 0 && (x + Y[i]) % 2 == 0) {
int k = (x + X[i]) / 2, j = (x + Y[i]) / 2;
if (i < j && k < H.size())
if (X[j] == x && X[k] == x) triples.insert({i, j, k});
}
}
}
}
return triples.size();
}
std::vector<int> construct_range(int M, int K) {
return {1, 1, 1};
}
