#include "triples.h"
#include <bits/stdc++.h>
using namespace std;
#define int long long
int ans = 0;
int n;
const int SQRT = 450;
void count1(vector<signed> &h){
for (int idx = 0; idx<n; idx++){
int x = h[idx];
int idh = idx - x;
if (idh < 0) continue;
int y = h[idh] - x;
if (y <= 0 || y==x) continue;
int idy = idx + y;
if (idy >= n) continue;
if (h[idy] == y) ans++;
}
}
void count2(vector<signed> &h){
for (int idx = 0; idx<n; idx++){
int x = h[idx];
int idy = idx + x;
if (idy >= n) continue;
int y = h[idy];
int idh = idx - y;
if (idh < 0) continue;
if (h[idh] == x + y) ans++;
}
}
void count3(vector<signed> &h){
for (int idx = 0; idx<n; idx++){
int x = h[idx];
int idh = idx + x;
if (idh >= n) continue;
int y = h[idh] - x;
if (y <= 0 || y==x) continue;
int idy = idh + y;
if (idy >= n) continue;
if (h[idy] == y) ans++;
}
}
void count4(vector<signed> &h){
vector<vector<int>> tl(3*n);
for (int i=0; i<n; i++){
int j = i + h[i]; if (j >= 3*n) continue;
tl[j].push_back(h[i]);
}
for (int idt = 0; idt<3*n; idt++){
if (tl[idt].size() > SQRT){
for (int idx = 0; idx < idt; idx++){
int x = h[idx];
int y = idt - idx; y -= x;
if (y<=0 || (y&1)) continue;
y/=2; int idy = idt - y;
int idh = idt - x - y;
if (h[idh] == x+y && h[idy] == y) ans++;
}
} else {
for (int y: tl[idt]){
for (int hv: tl[idt]){
int x = hv - y; if (x<=0) continue;
int idx = idt - hv - y;
if (idx <0) continue;
if (h[idx] == x) ans++;
}
}
}
}
}
long long count_triples(std::vector<signed> h) {
n = h.size(); ans = 0;
auto rh = h; reverse(rh.begin(), rh.end());
count1(h); count1(rh);
count2(h); count2(rh);
count3(h);
count4(h);
return ans;
}
std::vector<signed> construct_range(signed n, signed K) {
vector<signed> ans(2*n,0);
int st = sqrtl(n)-5; while (st*st < n) st++;
vector<int> S;
for (int i=0; i<st; i++) S.push_back(i);
for (int i=1; i<st; i++) S.push_back(i*st);
for (int nw: S){
for (int ot: S){
if (ot>=nw) continue;
if (ans[ot + nw]) continue;
ans[ot + nw] = nw - ot;
}
}
while ((int)ans.size()>n) ans.pop_back();
ans[0] = 1;
return ans;
}
