#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;
// 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_a, list_b;
vector<triple> unique_triples;
for (int i = N - 1; i >= 0; i--) {
int a = i + H[i];
list_a[a].push_back(i);
}
for (int j = 0; j < N; ++j) {
int a = j + H[j];
int b = H[j] - j;
list_a[a].pop_back();
//list_a[a].erase(find(list_a[a].begin(), list_a[a].end(), j));
for (int k : list_a[a])
for (int i : list_b[b + N])
if (j - i == H[k] && k - j == H[i] && k - i == H[j]) {
unique_triples.emplace_back(triple(i, j, k));
}
list_b[b + N].push_back(j);
}
//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};
}
