#include "triples.h"
#include <unordered_map>
#include <unordered_set>
long long count_easy_triples(const std::vector<int>& h) {
const int n = static_cast<int>(h.size());
std::vector<std::tuple<int, int, int>> triples;
for (int i = 0; i < n; i++) {
for (int j : {i - h[i], i + h[i]}) {
if (j < 0 || j >= n)
continue;
for (int k : {j - h[j], j + h[j]}) {
if (k < 0 || k >= n)
continue;
if (h[k] == abs(k - i)) {
triples.emplace_back(i, j, k);
}
}
for (int k : {i - h[j], i + h[j]}) {
if (k < 0 || k >= n)
continue;
if (h[k] == abs(k - j)) {
triples.emplace_back(i, j, k);
}
}
}
}
std::unordered_set<long long> deduped;
for (const auto& [i, j, k] : triples)
if (i != j && i != k && j != k) {
std::vector<int> vals{i, j, k};
std::sort(vals.begin(), vals.end());
deduped.insert(1ll * n * n * vals[0] + 1ll * n * vals[1] + vals[2]);
}
return deduped.size();
}
long long count_hard_triples(const std::vector<int>& h) {
const int n = static_cast<int>(h.size());
std::unordered_map<int, std::vector<int>> plus_groups;
std::unordered_map<int, std::vector<int>> minus_groups;
for (int i = 0; i < n; i++) {
plus_groups[i + h[i]].push_back(i);
minus_groups[i - h[i]].push_back(i);
}
long long ans = 0;
for (int j = 0; j < n; j++) {
const auto& i_cands = minus_groups[j - h[j]];
const auto& k_cands = plus_groups[j + h[j]];
const auto& small = i_cands.size() < k_cands.size() ? i_cands : k_cands;
for (const int i : small) {
const int k = (&small == &i_cands) ? j + h[i] : j - h[i];
if (k >= 0 && k < n && h[k] == std::abs(j - i)
&& h[i] != std::abs(j - i) && h[i] != std::abs(k - i)
&& h[j] != std::abs(j - i) && h[j] != std::abs(k - j)
&& h[k] != std::abs(k - i) && h[k] != std::abs(k - j)
&& i != j && i != k && j != k)
ans += (&small == &i_cands) ? (k + h[k] == j + h[j]) : (k - h[k] == j - h[j]);
}
}
return ans;
}
long long count_triples(std::vector<int> h) {
return count_easy_triples(h) + count_hard_triples(h);
}
std::vector<int> construct_range(int M, int K) {
return {0};
}
