#include "triples.h"
#include <set>
long long count_triples(std::vector<int> H)
{
int n = H.size();
std::vector<std::set<std::pair<int, int>>> s(n, std::set<std::pair<int, int>>());
for (int j = 1; j < n - 1; j++)
{
int i, k;
i = j - H[j];
if (i >= 0)
{
k = j + H[i];
if (k < n && k - i == H[k])
s[i].insert(std::pair<int, int>(j, k));
k = i + H[i];
if (k < n && k > j && H[k] == k - j)
s[i].insert(std::pair<int, int>(j, k));
}
k = j + H[j];
if (k < n)
{
i = j - H[k];
if (i >= 0 && k - i == H[i])
s[i].insert(std::pair<int, int>(j, k));
i = k - H[k];
if (i >= 0 && i < j && j - i == H[i])
s[i].insert(std::pair<int, int>(j, k));
}
}
std::vector<std::vector<int>> p(n, std::vector<int>());
for (int i = n - 1; i >= 0; i--)
{
int j;
if (i + H[i] < n)
{
for (int k : p[i + H[i]])
{
j = i + H[i];
if (j < k && H[j] == k - i)
s[i].insert(std::pair<int, int>(j, k));
j = i + H[k];
if (j < k && H[j] == k - i)
s[i].insert(std::pair<int, int>(j, k));
}
}
if (i - H[i] >= 0)
p[i - H[i]].push_back(i);
}
long long ans = 0;
for (int i = 0; i < s.size(); i++)
{
ans += s[i].size();
}
return ans;
}
std::vector<int> construct_range(int M, int K)
{
std::vector<int> ans(M, 1);
int x = M / 2;
for (int i = 0; i < x; i++)
{
ans[i] = x - i;
}
ans[x] = x + 1;
for (int i = x + 1; i < M; i++)
{
ans[i] = i - x;
}
return ans;
}
