#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;
vector<vector<int>> easy(vector<int> h) {
int n = sz(h);
vector<vector<int>> ret;
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];
if (k < n && h[k] == k - i) {
vector<int> a = {i, j, k};
sort(all(a));
ret.push_back(a);
}
};
auto case2 = [&]() { // to the left
int k = i + h[j];
if (k < n && k != j && h[k] == abs(k - j)) {
vector<int> a = {i, j, k};
sort(all(a));
ret.push_back(a);
}
};
auto case3 = [&]() {
int k = j - h[j];
if (k >= 0 && k != i && h[k] == k - i) {
vector<int> a = {i, j, k};
sort(all(a));
ret.push_back(a);
}
};
case1();
case2();
case3();
}
return ret;
}
vector<vector<int>> easy_rev(vector<int> h) {
int n = sz(h);
vector<vector<int>> ret;
for (int k = n - 1; k >= 0; --k) {
int j = k - h[k];
if (j < 0) continue;
auto case1 = [&]() {
int i = j - h[j];
if (i >= 0 && h[i] == k - j) {
vector<int> a = {i, j, k};
sort(all(a));
ret.push_back(a);
}
};
auto case2 = [&]() {
int i = k - h[j];
if (i >= 0 && i != j && h[i] == abs(i - j)) {
vector<int> a = {i, j, k};
sort(all(a));
ret.push_back(a);
}
};
auto case3 = [&]() {
int i = j + h[j];
if (i < n && i != k && h[i] == i - k) {
vector<int> a = {i, j, k};
sort(all(a));
ret.push_back(a);
}
};
case1();
case2();
case3();
}
return ret;
}
vector<vector<int>> small(int x, vector<int> group, vector<int>& h) {
int n = sz(h);
int m = sz(group);
vector<vector<int>> ret;
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];
if (k < n && h[k] == j - i && h[i] == k - j && h[j] == k - i) {
vector<int> a = {i, j, k};
sort(all(a));
ret.push_back(a);
}
};
auto case2 = [&]() {
int k = j - h[i];
if (k >= 0 && h[k] == j - i && h[i] == j - k && h[j] == abs(k - i)) {
vector<int> a = {i, j, k};
sort(all(a));
ret.push_back(a);
}
};
// cerr << "i and j: " << i << " " << j << "\n";
case1();
case2();
}
}
return ret;
}
vector<vector<int>> big(int x, vector<int> group, vector<int>& h) {
int n = sz(h);
int m = sz(group);
vector<vector<int>> ret;
sort(all(group));
auto check = [&](vector<int>& v) {
vector<int> a = {v[1] - v[0], v[2] - v[0], v[2] - v[1]};
vector<int> b = {h[v[0]], h[v[1]], h[v[2]]};
sort(all(a));
sort(all(b));
return a == b;
};
for (int k = 0; k < n; ++k) {
int z = k - h[k] - x;
if (z < 0 || z % 2) continue;
int i = z / 2;
int j = h[k] + i;
vector<int> v = {i, j, k};
sort(all(v));
if (check(v)) {
ret.push_back(v);
}
}
return ret;
}
vector<vector<int>> hard(vector<int> h) {
int n = sz(h);
map<int, vector<int>> groups;
for (int i = 0; i < n; ++i) {
groups[h[i] - i].push_back(i);
}
int sqr = ceil(sqrt(n));
vector<vector<int>> ret;
for (auto [x, group] : groups) {
// cerr << "x: " << x << "\n";
// cerr << "group: ";
// for (int e : group) {
// cerr << e << " ";
// }
// cerr << "\n";
if (sz(group) >= sqr) {
auto v = big(x, group, h);
for (auto a : v) ret.push_back(a);
} else {
auto v = small(x, group, h);
for (auto a : v) ret.push_back(a);
}
}
return ret;
}
long long count_triples(std::vector<int> h) {
int n = sz(h);
auto check = [&](vector<int>& v) {
vector<int> a = {v[1] - v[0], v[2] - v[0], v[2] - v[1]};
vector<int> b = {h[v[0]], h[v[1]], h[v[2]]};
sort(all(a));
sort(all(b));
return a == b;
};
auto a = easy(h);
auto b = easy_rev(h);
set<vector<int>> s;
for (auto v : a) {
if (check(v)) s.insert(v);
}
for (auto v : b) {
if (check(v)) s.insert(v);
}
// for (auto v : s) {
// for (int x : v) {
// cerr << x << " ";
// }
// cerr << "\n";
// }
ll cnt = sz(s);
// cerr << "cnt before hard: " << cnt << "\n";
auto c = hard(h);
for (auto v : c) {
if (check(v)) s.insert(v);
}
cnt = sz(s);
return cnt;
}
std::vector<int> construct_range(int M, int K) {
return {1, 1, 1};
}
