Problem p02691 (C++) — buggy submission, verdict: Time Limit Exceeded

#include <bits/stdc++.h>
using namespace std;
typedef int64_t ll;
typedef long double ld;
const ll MOD = 1000000007;
#define rep(i, n) for (ll i = 0; i < (ll)(n); i++)
long long gcd(long long a, long long b) {
  ll gcdmax = max(a, b);
  ll gcdmin = min(a, b);
  while (true) {
    if (gcdmax % gcdmin == 0)
      break;
    else
      gcdmax %= gcdmin;
    swap(gcdmin, gcdmax);
  }
  return gcdmin;
}
ll powerup(ll N, ll P, ll M) {
  if (P == 0)
    return 1;
  else if (P % 2 == 0) {
    ll t = powerup(N, P / 2, M);
    return t * t % M;
  } else
    return N * powerup(N, P - 1, M) % M;
}
vector<ll> find_divisor(ll N) {
  ll k = 1;
  while (k * k <= N) {
    k++;
  }
  vector<ll> A(1);
  rep(i, k) {
    if (i == 1)
      A.at(0) = 1;
    else if (i >= 2) {
      if (N % i == 0)
        A.push_back(i);
    }
  }
  ll t = 0;
  t = A.size();
  rep(i, t) {
    if (A.at(t - i - 1) * A.at(t - i - 1) != N)
      A.push_back(N / A.at(t - 1 - i));
  }
  return A;
}

bool isOK(ll index, ll key, vector<ll> A) {
  if (A.at(index) >= key)
    return true;
  else
    return false;
}

ll binary_search(ll key, vector<ll> A) {
  ll left = -1;
  ll right = A.size();
  while (right - left > 1) {
    int mid = left + (right - left) / 2;
    if (isOK(mid, key, A))
      right = mid;
    else
      left = mid;
  }
  return right;
}
bool isOKA(ll index, ll key, vector<ll> A) {
  if (A.at(index) > key)
    return true;
  else
    return false;
}

ll binary_searchA(ll key, vector<ll> A) {
  ll left = -1;
  ll right = A.size();
  while (right - left > 1) {
    int mid = left + (right - left) / 2;
    if (isOKA(mid, key, A))
      right = mid;
    else
      left = mid;
  }
  return right;
}
int main() {
  ll N;
  cin >> N;
  vector<ll> A(N);
  vector<ll> B(N);
  ll H;
  rep(i, N) {
    cin >> H;
    A.at(i) = H + i;
    B.at(i) = i - H;
  }
  vector<ll> C(N);
  vector<ll> D(N);
  rep(i, N) {
    if (A.at(i) < N)
      C.at(A.at(i))++;
    if (B.at(i) > 0)
      D.at(B.at(i))++;
  }
  ll ans = 0;
  rep(i, N) { ans += C.at(i) * D.at(i); }
  cout << ans << endl;
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
typedef int ll;
typedef long double ld;
const ll MOD = 1000000007;
#define rep(i, n) for (ll i = 0; i < (ll)(n); i++)
long long gcd(long long a, long long b) {
  ll gcdmax = max(a, b);
  ll gcdmin = min(a, b);
  while (true) {
    if (gcdmax % gcdmin == 0)
      break;
    else
      gcdmax %= gcdmin;
    swap(gcdmin, gcdmax);
  }
  return gcdmin;
}
ll powerup(ll N, ll P, ll M) {
  if (P == 0)
    return 1;
  else if (P % 2 == 0) {
    ll t = powerup(N, P / 2, M);
    return t * t % M;
  } else
    return N * powerup(N, P - 1, M) % M;
}
vector<ll> find_divisor(ll N) {
  ll k = 1;
  while (k * k <= N) {
    k++;
  }
  vector<ll> A(1);
  rep(i, k) {
    if (i == 1)
      A.at(0) = 1;
    else if (i >= 2) {
      if (N % i == 0)
        A.push_back(i);
    }
  }
  ll t = 0;
  t = A.size();
  rep(i, t) {
    if (A.at(t - i - 1) * A.at(t - i - 1) != N)
      A.push_back(N / A.at(t - 1 - i));
  }
  return A;
}

bool isOK(ll index, ll key, vector<ll> A) {
  if (A.at(index) >= key)
    return true;
  else
    return false;
}

ll binary_search(ll key, vector<ll> A) {
  ll left = -1;
  ll right = A.size();
  while (right - left > 1) {
    int mid = left + (right - left) / 2;
    if (isOK(mid, key, A))
      right = mid;
    else
      left = mid;
  }
  return right;
}
bool isOKA(ll index, ll key, vector<ll> A) {
  if (A.at(index) > key)
    return true;
  else
    return false;
}

ll binary_searchA(ll key, vector<ll> A) {
  ll left = -1;
  ll right = A.size();
  while (right - left > 1) {
    int mid = left + (right - left) / 2;
    if (isOKA(mid, key, A))
      right = mid;
    else
      left = mid;
  }
  return right;
}
int main() {
  ll N;
  cin >> N;
  vector<ll> A(N);
  vector<ll> B(N);
  ll H;
  rep(i, N) {
    cin >> H;
    A.at(i) = H + i;
    B.at(i) = i - H;
  }
  sort(B.begin(), B.end());
  ll ans = 0;
  rep(i, N) { ans += binary_searchA(A.at(i), B) - binary_search(A.at(i), B); }
  cout << ans << endl;
}