#include "triples.h"
#include <algorithm>
#include <vector>
using namespace std;
long long count_triples(std::vector<int> H) {
int N = H.size();
long long ans = 0;
// 情况 B1
for (int i = 0; i < N; ++i) {
int j = i + H[i];
if (j < N) {
int k = j + H[j];
if (k < N && H[k] == k - i) ans++;
}
}
// 情况 B2
for (int j = 0; j < N; ++j) {
int i = j - H[j];
if (i >= 0) {
int k = j + H[i];
if (k < N && H[k] == k - i) ans++;
}
}
// 容斥重叠 B12
for (int i = 0; i < N; ++i) {
int j = i + H[i];
if (j < N) {
int k = j + H[i];
if (k < N && H[j] == H[i] && H[k] == k - i) ans--;
}
}
// 情况 A1
for (int i = 0; i < N; ++i) {
int j = i + H[i];
if (j < N) {
int k = i + H[j];
if (k < N && j < k && H[k] == k - j) ans++;
}
}
// 情况 A2 (使用启发式小集合暴力，保证 O(N * sqrt(N)))
int OFFSET = N + 5;
vector<vector<int>> bucketD(2 * N + 10);
vector<vector<int>> bucketV(2 * N + 10);
for (int x = 0; x < N; ++x) {
bucketD[x - H[x] + OFFSET].push_back(x);
bucketV[x + H[x]].push_back(x);
}
for (int j = 0; j < N; ++j) {
int U = j - H[j] + OFFSET;
int V = j + H[j];
if (U < 0 || U >= 2 * N + 10 || V < 0 || V >= 2 * N + 10) continue;
int countD = lower_bound(bucketD[U].begin(), bucketD[U].end(), j) - bucketD[U].begin();
int countV = bucketV[V].end() - upper_bound(bucketV[V].begin(), bucketV[V].end(), j);
if (countD <= countV) {
for (int idx = 0; idx < countD; ++idx) {
int i = bucketD[U][idx];
int pos = i + H[j];
if (pos < N && pos > j && pos + H[pos] == V) {
ans++;
}
}
} else {
auto it = upper_bound(bucketV[V].begin(), bucketV[V].end(), j);
for (; it != bucketV[V].end(); ++it) {
int pos = *it;
int i = pos - H[j];
if (i >= 0 && i < j && i - H[i] + OFFSET == U) {
ans++;
}
}
}
}
// 容斥重叠 A12
for (int i = 0; i < N; ++i) {
int j = i + H[i];
if (j < N) {
int k = j + H[i];
if (k < N && H[j] == k - i && H[k] == H[i]) ans--;
}
}
// 情况 C1
for (int i = 0; i < N; ++i) {
int k = i + H[i];
if (k < N) {
int j = k - H[k];
if (j > i && j < k && j - i == H[j]) ans++;
}
}
// 情况 C2
for (int j = 0; j < N; ++j) {
int k = j + H[j];
if (k < N) {
int i = j - H[k];
if (i >= 0 && i < j && H[i] == k - i) ans++;
}
}
// 容斥重叠 C12
for (int j = 0; j < N; ++j) {
int k = j + H[j];
if (k < N) {
int i = j - H[j];
if (i >= 0 && H[k] == H[j] && H[i] == k - i) ans--;
}
}
return ans;
}
std::vector<int> construct_range(int M, int K) {
return {1, 1, 1};
}
