# | Time | Username | Problem | Language | Result | Execution time | Memory |
---|---|---|---|---|---|---|---|
1250054 | testacc11 | Triple Peaks (IOI25_triples) | C++20 | 0 ms | 0 KiB |
#include <bits/stdc++.h>
using namespace std;
using ll = long long;
ll count_triples(const vector<int>& H) {
int N = H.size();
// L1[j] = list of all i<j with H[i] = j-i
vector<vector<int>> L1(N);
// S[T] = list of all k with k+H[k] = T
vector<vector<int>> S(2*N);
for(int i = 0; i < N; i++) {
int j1 = i + H[i];
if(j1 < N) L1[j1].push_back(i);
int T = i + H[i];
S[T].push_back(i);
}
ll ans = 0;
for(int j = 0; j < N; j++) {
int hj = H[j];
// Case A: H[j] = d1 = j-i
{
int i = j - hj;
if(i >= 0) {
int hi = H[i];
int k = j + hi;
if(k < N && hi + hj == H[k]) {
ans++;
}
}
}
// Case B: H[j] = d2 = k-j
{
int k = j + hj;
if(k < N) {
int hk = H[k];
int i = k - hk;
if(i >= 0 && i < j && hk - hj == H[i]) {
ans++;
}
}
}
// Case C1: H[j] = d3, with H[i]=d1 and H[k]=d2
for(int i : L1[j]) {
int hi = H[i];
int k = i + hj;
if(k > j && k < N && H[k] == hj - hi) {
ans++;
}
}
// Case C2: H[j] = d3, with H[k]=d1 and H[i]=d2
int Tj = j + hj;
if(Tj < (int)S.size()) {
for(int k : S[Tj]) {
if(k <= j) continue;
int i = k - hj;
if(i >= 0 && i < j && H[i] == k - j) {
ans++;
}
}
}
}
return ans;
}
int main(){
ios::sync_with_stdio(false);
cin.tie(nullptr);
int part;
cin >> part;
if(part != 1) return 0; // only Part I
int N;
cin >> N;
vector<int> H(N);
for(int i = 0; i < N; i++) {
cin >> H[i];
}
ll result = count_triples(H);
cout << result << "\n";
return 0;
}