#include "triples.h"
#include <vector>
using namespace std;
// Part I: O(N) count by handling the three ways the two smallest distances
// can come from two of the heights.
long long count_triples(vector<int> H) {
int N = H.size();
long long ans = 0;
// Case A: the two *smallest* distances are d(i,j) = H[i] and d(j,k) = H[j].
// ⇒ j = i + H[i], k = j + H[j], and H[k] must equal H[i] + H[j].
for(int i = 0; i < N; i++){
int j = i + H[i];
if(j < N){
int k = j + H[j];
if(k < N && H[k] == H[i] + H[j]) {
ans++;
}
}
}
// Case B: the two *smallest* distances are d(i,j) = H[i] and d(i,k) = H[k].
// ⇒ j = k - H[k], k = i + H[i], and H[j] + H[k] == H[i] + H[j]? Actually:
// we want {H[i],H[j],H[k]} = {H[i],H[k], H[i]+H[k]}, so H[j] == H[i] + H[k].
for(int i = 0; i < N; i++){
int k = i + H[i];
if(k < N){
int j = k - H[k];
if(j > i && j < N && H[j] == H[i] + H[k]) {
ans++;
}
}
}
// Case C: the two *smallest* distances are d(i,k) = H[k] and d(j,k) = H[j].
// ⇒ i = j - H[j], k = j + H[j], and H[i] + H[k] == H[j].
for(int j = 0; j < N; j++){
int i = j - H[j];
int k = j + H[j];
if(i >= 0 && k < N && H[i] + H[k] == H[j]) {
ans++;
}
}
return ans;
}
// Part II: pack as many [1,1,2] blocks as will fit in M:
// each block is size 3 and contributes exactly one mythical triple.
vector<int> construct_range(int M, int K) {
int blocks = M / 3; // maximum disjoint [1,1,2] blocks
int want = min(blocks, K); // if K smaller, only need K blocks
int N = want * 3;
if(N < 3) {
// must output at least 3 peaks
return vector<int>{1,1,2};
}
vector<int> H;
H.reserve(N);
for(int b = 0; b < want; b++){
H.push_back(1);
H.push_back(1);
H.push_back(2);
}
return H;
}
