#include <iostream>
#include <vector>
#include <numeric>
#include <map>
#include <set>
#include <tuple>
#include <algorithm>
// Placeholder for Part II, required by the grader to compile.
std::vector<int> construct_range(int M, int K) {
return {};
}
// Correctly counts mythical triples without overcounting.
long long count_triples(std::vector<int> H) {
int N = H.size();
if (N < 3) {
return 0;
}
// Use a set of tuples to store unique mythical triples found.
// This elegantly solves the overcounting problem.
std::set<std::tuple<int, int, int>> found_triples;
// We check all 6 permutations. For each valid triple (i, j, k) found,
// we insert it into the set. The set ensures uniqueness.
// Iterate through the middle element `j` to structure the search
for (int j = 1; j < N - 1; ++j) {
// --- Find potential i's (i < j) ---
// Case where H[j] = j - i => i = j - H[j]
if (H[j] > 0 && j - H[j] >= 0) {
int i = j - H[j];
if (i < j) {
// Now we have (i, j) and need to find k > j
// Case 3: H[i]=b, H[j]=a, H[k]=a+b
int k = j + H[i];
if (k > j && k < N && H[k] == H[i] + H[j]) {
found_triples.insert({i, j, k});
}
// Case 5: H[i]=a+b, H[j]=a, H[k]=b
k = i + H[i];
if (k > j && k < N && H[i] > H[j] && H[k] == H[i] - H[j]) {
found_triples.insert({i, j, k});
}
}
}
}
// --- Forward pass for cases anchored at i ---
for (int i = 0; i < N; ++i) {
// Case where H[i] = j - i => j = i + H[i]
if (H[i] > 0 && i + H[i] < N) {
int j = i + H[i];
if (j > i) {
// Now we have (i, j) and need to find k > j
// Case 1: H[i]=a, H[j]=b, H[k]=a+b
if (H[j] > 0) {
int k = j + H[j];
if (k > j && k < N && H[k] == H[i] + H[j]) {
found_triples.insert({i, j, k});
}
}
// Case 2: H[i]=a, H[j]=a+b, H[k]=b
int k = i + H[j];
if (k > j && k < N && H[j] > H[i] && H[k] == H[j] - H[i]) {
found_triples.insert({i, j, k});
}
}
}
}
// --- Backward pass for cases anchored at k ---
for (int k = N-1; k >= 0; --k) {
// Case where H[k] = k - j => j = k - H[k]
if(H[k] > 0 && k - H[k] >= 0) {
int j = k - H[k];
if (j < k) {
// Now we have (j, k) and need to find i < j
// Case 6: H[i]=a+b, H[j]=b, H[k]=a
int i = k - H[i]; // This is not right, H[i] is unknown.
// Let's re-derive: k-i = H[i], so k = i+H[i]
// This case is difficult to anchor at k.
}
}
}
// A more robust map-based approach for the remaining complex cases
// to ensure all are found without overcounting.
std::map<int, std::vector<int>> i_plus_Hi;
std::map<int, std::vector<int>> k_minus_Hk;
for(int i = 0; i < N; ++i) i_plus_Hi[i+H[i]].push_back(i);
for(int k = 0; k < N; ++k) k_minus_Hk[k-H[k]].push_back(k);
// Case 4: H[i]=b, H[j]=a+b, H[k]=a. => k-i=H[j], j-i=H[k], k-j=H[i]
// From this: j = i + H[k] and k = j + H[i] => k-H[i] = i+H[k]
for(auto const& [val, i_list] : i_plus_Hi) {
if(k_minus_Hk.count(val)) {
for(int i : i_list) {
for(int k : k_minus_Hk[val]) {
if (i < k) {
int j = i + H[k];
if (j > i && j < k && H[j] == H[i] + H[k]) {
found_triples.insert({i, j, k});
}
}
}
}
}
}
// Case 6: H[i]=a+b, H[j]=b, H[k]=a => k-i=H[i], j-i=H[k], k-j=H[j]
// From this: k=i+H[i] and j=i+H[k]
// And k-H[j]=j => k=j+H[j]
// So we need i+H[i] = j+H[j] = k
for(auto const& [k, i_list] : i_plus_Hi) {
if (k >= 0 && k < N) { // k must be a valid index
for(int i : i_list) {
// Now find a `j` such that j+H[j] = k and i < j < k
if(i_plus_Hi.count(k)) {
for(int j : i_plus_Hi[k]) {
if (i < j && j < k) { // Ensure order
if (H[i] > H[j] && H[k] == H[i] - H[j]) {
found_triples.insert({i, j, k});
}
}
}
}
}
}
}
return found_triples.size();
}
