#include "triples.h"
#include <bits/stdc++.h>
using namespace std;
set<pair<int, int>> res[202020];
vector<int> pts[404040];
unordered_map<int, unordered_set<int>> nums;
map<pair<int, int>, long long> memo;
long long count_triples(vector<int> H) {
int N = H.size();
long long ans = 0;
int i, j, k;
for (j = 0; j < N; j++) {
// case 1
i = j - H[j];
if (i >= 0) {
k = i + H[i];
if (k < N && H[k] == k-j) res[j].insert({i, k});
k = j + H[i];
if (k < N && H[k] == k-i) res[j].insert({i, k});
}
// case 2
k = j + H[j];
if (k < N) {
i = k - H[k];
if (i >= 0 && H[i] == j-i) res[j].insert({i, k});
i = j - H[k];
if (i >= 0 && H[i] == k-i) res[j].insert({i, k});
}
}
// case 3-1
for (i = 0; i < N; i++) {
j = i + H[i];
if (j >= N) continue;
k = i + H[j];
if (k >= N) continue;
if (H[k] == k-j) res[j].insert({i, k});
}
// case 3-2
int B = pow(N, 2.0/3.0);
for (int i = 0; i < N; i++) {
nums[i - H[i]].insert(i + H[i]);
}
for (int i = 0; i < N; i++) {
int a = i - H[i], b = i + H[i];
if (nums[a].size() > B && nums[b].size() > B) {
if (memo.find({a, b}) != memo.end()) ans += memo[{a, b}];
else {
for (auto &n : nums[a]) {
if (nums[b].find(n) != nums[b].end()) {
memo[{a, b}]++;
}
}
ans += memo[{a, b}];
}
}
else {
if (nums[a].size() > nums[b].size()) {
swap(a, b);
}
for (auto &n : nums[a]) {
if (nums[b].find(n) != nums[b].end()) {
ans++;
}
}
}
}
for (j = 0; j < N; j++) {
for (auto &p : res[j]) {
i = p.first;
k = p.second;
if (j - i == H[k] && k - j == H[i] && k - i == H[j]) ans--;
}
ans += res[j].size();
}
return ans;
}
vector<int> construct_range(int M, int K) {
int N = M;
vector<int> result(N);
result[0] = 4;
result[1] = 1;
for (int i = 2; i < N; i++) result[i] = i-1;
for (int i = 2; i < N; i += 2) swap(result[i], result[i+1]);
return result;
}
