#include "triples.h"
#include <bits/stdc++.h>
using namespace std;
bool checkTriple(int i, int j, int k, vector<int>&arr){
if(!(i<j&&j<k)){
return 0;
}
int n = arr.size();
if(i<0||j<0||k<0||i>=n||j>=n||k>=n){
return 0;
}
vector<int>arr1 = {j-i,k-i,k-j};
vector<int>arr2 = {arr[i],arr[j],arr[k]};
sort(arr1.begin(),arr1.end());
sort(arr2.begin(),arr2.end());
if(arr1[0]==arr2[0]&&arr1[1]==arr2[1]&&arr1[2]==arr2[2]){
return 1;
}
return 0;
}
long long count_triples(vector<int> arr) {
int n = arr.size();
set<array<int,3>>triples;
unordered_map<int,vector<int>>mp;
for(int i = 0;i<n;i++){
mp[arr[i]-i].push_back(i);
}
for(int i = 0;i<n;i++){
int a,b,c;
//case 1:
a = i;
b = arr[a]+a;
if(b>=0&&b<n)
c = arr[b]+a;
if(checkTriple(a,b,c,arr)){
triples.insert({a,b,c});
}
//case 2:
a=i;
c=arr[a]+a;
if(c>=0&&c<n)
b=c-arr[c];
if(checkTriple(a,b,c,arr)){
triples.insert({a,b,c});
}
//case 3:
a=i;
c=arr[a]+a;
if(c>=0&&c<n)
b=arr[c]+a;
if(checkTriple(a,b,c,arr)){
triples.insert({a,b,c});
}
//case 4:
a=i;
b=arr[a]+a;
if(b>=0&&b<n)
c=arr[b]+b;
if(checkTriple(a,b,c,arr)){
triples.insert({a,b,c});
}
//case 5:
//later
//case 6:
b=i;
a=b-arr[b];
if(a>=0&&a<n)
c=arr[a]+b;
if(checkTriple(a,b,c,arr)){
triples.insert({a,b,c});
}
}
int lim = sqrt(n);
for(pair<int,vector<int>>p:mp){
int siz = p.second.size();
if(siz>lim)
continue;
for(int i = 0;i<siz;i++){
for(int j = i+1;j<siz;j++){
int a = p.second[i];
int b = p.second[j];
int c = arr[b]+a;
if(checkTriple(a,b,c,arr)){
triples.insert({a,b,c});
}
}
}
}
for(pair<int,vector<int>>p:mp){
int siz = p.second.size();
if(siz<=lim)
continue;
int x = p.first;
for(int c = 0;c<n;c++){
int a = (c-x-arr[c])/2;
int b = (arr[c]+c-x)/2;
if(checkTriple(a,b,c,arr)){
triples.insert({a,b,c});
}
}
}
return triples.size();
}
vector<int> construct_range(int n, int K) {
vector<int>arr(n);
arr[0]=1;
iota(arr.begin()+1,arr.end(),1);
return arr;
}
