#include "triples.h"
#include <bits/stdc++.h>
using namespace std;
const int N = 2e5 + 5;
const int K = 650;
long long calcTriangle(vector<int> h)
{
int n = size(h);
vector<vector<int>> a(n * 3);
vector<pair<int, int>> edges;
for (int i = 0; i < n * 3; i++)
a[i].clear();
for (int i = 0; i < n; i++)
{
int l = i - h[i], r = i + h[i];
a[l + n].push_back(r + n);
a[r + n].push_back(l + n);
edges.push_back({l + n, r + n});
}
vector<int> bigId(n * 3, -1);
int bigCnt = 0;
for (int i = 0; i < n * 3; i++)
if (size(a[i]) > K)
bigId[i] = bigCnt++;
vector<vector<int>> adj(bigCnt, vector<int>(bigCnt));
for (auto [x, y] : edges)
if (bigId[x] >= 0 && bigId[y] >= 0)
adj[bigId[x]][bigId[y]] = adj[bigId[y]][bigId[x]] = 1;
long long res = 0;
vector<int> flag(n * 3, -1);
for (int x = 0; x < n * 3; x++)
{
for (int y : a[x])
flag[y] = x;
vector<int> bigs;
for (int y : a[x])
if (bigId[y] >= 0) bigs.push_back(bigId[y]);
else
{
for (int z : a[y])
if (flag[z] == x && y < z)
res++;
}
for (int i = 0; i < size(bigs); i++)
for (int j = i + 1; j < size(bigs); j++)
res += adj[bigs[i]][bigs[j]];
}
assert(res % 3 == 0);
return res / 3;
}
long long count_triples(vector<int> h)
{
int n = size(h);
vector<int> l(n), r(n);
for (int i = 0; i < n; i++)
{
l[i] = i - h[i];
r[i] = i + h[i];
}
long long res = 0;
for (int i = 0; i < n; i++)
{
int j = r[i];
if (j < n)
{
// 2.3..5
res += r[j] < n && h[r[j]] == h[i] + h[j];
// 2.5..3
if (h[j] > h[i])
{
int k = j + h[j] - h[i];
if (k < n && l[k] == j)
res++;
}
// 5.3..2
if (h[j] < h[i])
{
int k = i + h[j];
if (r[k] == j && h[k] != h[j]) // exclude 211
res++;
}
}
}
for (int i = 0; i < n; i++)
{
int j = l[i];
if (j < 0)
continue;
// 5.2..3
if (l[j] >= 0 && h[l[j]] == h[i] + h[j])
res++;
// 3.2..5
int k = i + h[j];
if (r[j] != i && k < n && h[k] == h[i] + h[j]) // exclude 112
res++;
}
// k.2.i5.3..j
res += calcTriangle(h);
// exclude 121
for (int i = 0; i < n; i++)
if (h[i] % 2 == 0)
{
int u = i - h[i] / 2, v = i + h[i] / 2;
if (u >= 0 && h[u] == h[i] / 2 && v < n && h[v] == h[i] / 2)
res--;
}
return res;
}
vector<int> construct_range(int M, int K) {
vector<int> PATTERN = {2, 3, 4, 1, 2, 1, 3};
vector<int> ans(M);
for (int i = 0; i < M; i++)
ans[i] = PATTERN[i % size(PATTERN)];
return ans;
}
