#include "triples.h"
#include <bits/stdc++.h>
#define all(x) begin(x), end(x)
#define sz(x) (int)(x).size()
using namespace std;
using ll = long long;
long long easy(vector<int> h) {
ll cnt = 0;
int n = sz(h);
for (int i = 0; i < n; ++i) {
int j = i + h[i];
if (j >= n) continue;
auto case1 = [&]() { // to the right
int k = j + h[j];
return (int)(k < n && h[k] == k - i);
};
auto case2 = [&]() { // to the left
int k = i + h[j];
return (int)(k < n && k != j && h[k] == abs(k - j));
};
// cerr << "i: " << i << "\n";
// cerr << "case1: " << case1() << "\n";
// cerr << "cas2: " << case2() << "\n";
cnt += case1();
cnt += case2();
}
return cnt;
}
ll small(int x, vector<int> group, vector<int>& h) {
int n = sz(h);
int m = sz(group);
ll cnt = 0;
sort(all(group));
// cerr << "x: " << x << "\n";
for (int u = 0; u < m - 1; ++u) {
int i = group[u];
for (int z = u + 1; z < m; ++z) {
int j = group[z];
auto case1 = [&]() {
int k = j + h[i];
return (int)(k < n && h[k] == j - i && h[i] == k - j && h[j] == k - i);
};
auto case2 = [&]() {
int k = j - h[i];
return (int)(k >= 0 && h[k] == j - i && h[i] == j - k && h[j] == abs(k - i));
};
// cerr << "i and j: " << i << " " << j << "\n";
cnt += case1() + case2();
}
}
return cnt;
}
long long hard(vector<int> h) {
int n = sz(h);
map<int, vector<int>> groups;
ll cnt = 0;
for (int i = 0; i < n; ++i) {
groups[h[i] - i].push_back(i);
}
int sqr = ceil(sqrt(n));
for (auto [x, group] : groups) {
// cerr << "x: " << x << "\n";
// cerr << "group: ";
// for (int e : group) {
// cerr << e << " ";
// }
// cerr << "\n";
if (sz(group) >= sqr) {
// cnt += big(x, group);
cnt += small(x, group, h);
} else {
cnt += small(x, group, h);
}
}
return cnt;
}
long long count_triples(std::vector<int> h) {
int n = sz(h);
ll cnt = easy(h);
reverse(all(h));
cnt += easy(h);
reverse(all(h));
cnt += hard(h);
return cnt;
}
std::vector<int> construct_range(int M, int K) {
return {1, 1, 1};
}
