#include "triples.h"
#include <functional>
#include <map>
#include <assert.h>
using namespace std;
struct triple {
int a, b, c;
triple(int x, int y, int z) {
if (x > y) swap(x, y);
if (y > z) swap(y, z);
if (x > y) swap(x, y);
a = x;
b = y;
c = z;
}
bool operator<(const triple& other) const {
if (a != other.a) return a < other.a;
if (b != other.b) return b < other.b;
return c < other.c;
}
bool operator==(const triple& other) const {
return a == other.a && b == other.b && c == other.c;
}
};
long long count_triples(std::vector<int> H) {
int N = H.size();
long long int result = 0;
int sq = sqrt(N) + 1;
// vector<long long> count_a(2*N + 1, 0);
// vector<long long> count_b(2*N + 1, 0);
// for (int i = 0; i < N; ++i) {
// int a = i + H[i];
// count_a[a]++;
// }
// for (int i = 0; i < N; ++i) {
// int a = i + H[i];
// int b = H[i] - i;
// count_a[a]--;
// result += count_a[a] * count_b[b + N];
// count_b[b + N]++;
// }
map<int, vector<int>> list_x;
vector<triple> unique_triples;
for (int i = 0; i < N; i++) {
int x = H[i] - i;
list_x[x + N].push_back(i);
}
for (auto [x, indices] : list_x) {
int size = indices.size();
std::function<void(int, int, int)> check_add_triple = [&](int i, int j, int k) {
if (i < 0 || j < 0 || k < 0 || i >= N || j >= N || k >= N) return;
if (i - j == H[k] && j - k == H[i] && i - k == H[j]) {
unique_triples.emplace_back(triple(i, j, k));
}
};
if (size < sq) {
for (int ind1 = 0; ind1 < size; ++ind1) {
for (int ind2 = ind1 + 1; ind2 < size; ++ind2) {
int j = indices[ind1];
int k = indices[ind2];
if (j < k)
swap(j, k);
int i = H[j] + k;
check_add_triple(i, j, k);
i = H[k] + j;
check_add_triple(i, j, k);
}
}
}
else {
for (int i = 0; i < N; ++i) {
int k = x + H[i] - i;
if (k % 2 != 0) continue;
k /= 2;
int j = H[i] + k;
check_add_triple(i, j, k);
j = k - H[i];
check_add_triple(i, j, k);
}
}
}
//assert(result % 6 == 0);
//result /= 6;
for (int i = 0; i < N; ++i) {
std::function<void(int)> check_j = [&](int j) {
if (j < N && j >= 0) {
int k = H[j] + j;
if (k < N && (k + H[k] == i || k - H[k] == i)) {
unique_triples.emplace_back(triple(i, j, k));
}
k = j - H[j];
if (k >= 0 && (k + H[k] == i || k - H[k] == i)) {
unique_triples.emplace_back(triple(i, j, k));
}
k = H[j] + i;
if (k < N && (k + H[k] == j || k - H[k] == j)) {
unique_triples.emplace_back(triple(i, j, k));
}
k = i - H[j];
if (k >= 0 && (k + H[k] == j || k - H[k] == j)) {
unique_triples.emplace_back(triple(i, j, k));
}
}
};
check_j(i + H[i]);
check_j(i - H[i]);
}
sort(unique_triples.begin(), unique_triples.end());
unique_triples.erase(unique(unique_triples.begin(), unique_triples.end()), unique_triples.end());
// for (triple& t : unique_triples) {
// int dist[] = {
// abs(t.a - t.b),
// abs(t.b - t.c),
// abs(t.c - t.a)
// };
// int hs[] = {
// H[t.a],
// H[t.b],
// H[t.c]
// };
// sort(dist, dist + 3);
// sort(hs, hs + 3);
// //if (dist[0] == hs[0] && dist[1] == hs[1] && dist[2] == hs[2]) {
// result++;
// //}
// }
return unique_triples.size();
}
std::vector<int> construct_range(int M, int K) {
return {1, 1, 1};
}
