#include "triples.h" #include<bits/stdc++.h> using namespace std; typedef long long ll; long long count_triples(std::vector<int> H) { int N=H.size(), OFFSET=H.size(), ans=0; vector<vector<int>> sum(2*N+1, vector<int>()), diff(2*N+1, vector<int>()); for (int i = 0; i < N; i++) { sum[i+H[i]].push_back(i); diff[i-H[i]+OFFSET].push_back(i); } // LET i>j>k for (int t = 0; t < N; t++) { { // h_i=i-k int i=t; int k=i-H[i]; if (k<0 || H[k]>=H[i]) goto next; { // h_k=i-j int j=i-H[k]; if (H[j]==j-k) ans++; } { // h_k=j-k -- prevent double counting int j=k+H[k]; if (H[j]==i-j && i-j!=j-k) ans++; } } next: { // h_k=i-k int k=t; int i=H[k]+k; if (i>=N || H[i]>=H[k]) goto end; { // h_i=i-j int j=i-H[i]; if (H[j]==j-k) ans++; } { // h_i=j-k -- double counting int j=H[i]+k; if (H[j]==i-j && i-j!=j-k) ans++; } } end: continue; } // h_j=i-k for (int j = 0; j < N; j++) { // h_k=j-k -- double counting for next part for (int t = 0; t < sum[j].size(); t++) { int k=sum[j][t]; if (k>=j) break; if (H[k]>=H[j]) continue; int Hi=H[j]-H[k]; int i=Hi+j; if (i<N && H[i]==Hi && i-j!=j-k) ans++; } int sum_start=upper_bound(sum[j+H[j]].begin(), sum[j+H[j]].end(), j)-sum[j+H[j]].begin(); int diff_end=lower_bound(diff[j-H[j]+OFFSET].begin(), diff[j-H[j]+OFFSET].end(), j)-diff[j-H[j]+OFFSET].begin(); if (sum[j+H[j]].size()-sum_start < diff_end) { // start with h_i+i=h_j+j auto ref=sum[j+H[j]]; for (int t = sum_start; t < ref.size(); t++) { int i=ref[t]; int k=j-H[i]; if (k<0) continue; if (H[k]+H[i]==H[j]) ans++; } } else { // start with j-h_j=k-h_k auto ref=diff[j-H[j]+OFFSET]; for (int t = 0; t < diff_end; t++) { int k=ref[t]; int i=H[k]+j; if (i>=N) continue; if (H[k]+H[i]==H[j]) ans++; } } } return ans; } std::vector<int> construct_range(int M, int K) { return {1, 1, 1}; }