#include <iostream>
#include <vector>
#include <numeric>
#include <map>
// Placeholder function for Part II to allow the program to compile.
// The grader requires this function to exist.
std::vector<int> construct_range(int M, int K) {
// This is a dummy implementation.
// For a real submission for Part II, this would need to be implemented.
return {1, 2, 3};
}
// Function for Part I, renamed to match the grader's expectation.
long long count_triples(std::vector<int> H) {
long long count = 0;
int N = H.size();
// We analyze all 6 permutations of {H[i], H[j], H[k]} vs {a, b, a+b}
// where a = j-i, b = k-j, a+b = k-i.
// We must ensure 0 <= i < j < k < N for all triples.
// Case 1: H[i]=a, H[j]=b, H[k]=a+b
for (int i = 0; i < N; ++i) {
if (H[i] > 0) {
int j = i + H[i];
if (j < N && H[j] > 0) {
int k = j + H[j];
if (k < N && H[k] == H[i] + H[j]) {
count++;
}
}
}
}
// Case 2: H[i]=a, H[j]=a+b, H[k]=b
for (int i = 0; i < N; ++i) {
if (H[i] > 0) {
int j = i + H[i];
if (j < N && H[j] > H[i]) {
int k = i + H[j];
if (k > j && k < N && H[k] == H[j] - H[i]) {
count++;
}
}
}
}
// Case 3: H[i]=b, H[j]=a, H[k]=a+b
std::map<int, std::vector<int>> j_minus_Hj;
for(int j = 0; j < N; ++j) {
if (j - H[j] >= 0) {
j_minus_Hj[j - H[j]].push_back(j);
}
}
for(int i = 0; i < N; ++i) {
if(j_minus_Hj.count(i)) {
for(int j : j_minus_Hj[i]) {
if(i < j) {
int k = j + H[i];
if(k > j && k < N && H[k] == H[i] + H[j]) {
count++;
}
}
}
}
}
// Case 4: H[i]=b, H[j]=a+b, H[k]=a
std::map<int, std::vector<int>> k_minus_Hk;
for (int k = 0; k < N; ++k) {
k_minus_Hk[k - H[k]].push_back(k);
}
for (int i = 0; i < N; ++i) {
int target = i + H[i];
if (k_minus_Hk.count(target)) {
for (int k : k_minus_Hk[target]) {
if (i < k) {
int j = k - H[i];
if (i < j && j < k && H[j] > H[i] && H[j] > H[k] && H[j] == H[i] + H[k]) {
count++;
}
}
}
}
}
// Case 5: H[i]=a+b, H[j]=a, H[k]=b
std::map<int, std::vector<int>> i_plus_Hi;
for (int i = 0; i < N; i++) {
if (H[i] > 0) {
int k = i + H[i];
if (k < N) {
if (i_plus_Hi.count(k)) {
for (int j : i_plus_Hi[k]) {
if (i < j && H[j] > 0 && H[k] > 0 && H[i] == H[j] + H[k]) {
count++;
}
}
}
}
}
// This map is used for a different case, but we can populate it here
int j_minus_Hj_val = i - H[i];
if(j_minus_Hj_val >= 0) i_plus_Hi[j_minus_Hj_val].push_back(i);
}
// Case 6: H[i]=a+b, H[j]=b, H[k]=a
std::map<int, std::vector<int>> i_plus_Hi_map;
for(int j=0; j<N; ++j) {
int val = j + H[j];
if(i_plus_Hi_map.count(val)) {
for(int i : i_plus_Hi_map[val]) {
if (i < j) {
int k = i + H[i];
if(k > j && k < N) {
if (H[i] > H[j] && H[k] == H[i] - H[j]) {
count++;
}
}
}
}
}
i_plus_Hi_map[j + H[j]].push_back(j);
}
return count;
}
