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

#include <bits/stdc++.h>
using namespace std;
#define lli long long int
#define rep(i, n) for (lli i = 0; i < n; i++)
#define repb(i, n) for (lli i = n - 1; i >= 0; i--)
#define pb push_back
#define mp make_pair
#define bg() begin()
#define en() end()
#define ft first
#define sc second
#define vect_lli_it vector<lli>::iterator
#define set_lli_it set<lli>::iterator
#define inf 1e18
#define all(X) X.begin(), X.end()
#define iterate(X) for (it = X.begin(); it != X.end(); it++)
lli combi(lli n, lli k) {
  lli ans = 1;
  k = k > n - k ? n - k : k;
  lli j = 1;
  for (; j <= k; j++, n--) {
    if (n % j == 0) {
      ans *= n / j;
    } else if (ans % j == 0) {
      ans = ans / j * n;
    } else {
      ans = (ans * n) / j;
    }
  }
  return ans;
}
lli gcd(lli a, lli b) {
  if (b == 0)
    return a;
  return gcd(b, a % b);
}
lli findlcm(vector<lli> arr, lli n) {
  lli ans = arr[0];
  for (lli i = 1; i < n; i++)
    ans = (((arr[i] * ans)) / (gcd(arr[i], ans)));

  return ans;
}
lli power(lli x, lli y, lli p) {
  lli res = 1;
  x = x % p;
  if (x == 0)
    return 0;
  while (y > 0) {
    if (y & 1)
      res = (res * x) % p;
    y = y >> 1;
    x = (x * x) % p;
  }
  return res;
}
int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(NULL);
  lli n;
  cin >> n;
  vector<lli> arr;
  lli x;
  rep(i, n) {
    cin >> x;
    arr.pb(x);
  }
  vector<lli> sum_arr;
  rep(i, n) { sum_arr.pb(i + arr[i] + 1); }
  vector<lli> diff_arr;
  rep(i, n) { diff_arr.pb(i - arr[i] + 1); }
  sort(all(sum_arr));
  sort(all(diff_arr));
  vect_lli_it itl, itu;
  lli ans = 0;
  rep(i, sum_arr.size()) {
    itl = lower_bound(all(diff_arr), sum_arr[i]);
    itu = upper_bound(all(diff_arr), sum_arr[i]);
    if (itl != diff_arr.en() && itu != diff_arr.en() && *itl == sum_arr[i]) {
      ans = ans + (itu - diff_arr.bg()) - (itl - diff_arr.bg());
    } else if (itl != diff_arr.en() && *itl == sum_arr[i]) {
      ans = ans + diff_arr.size() - (itl - diff_arr.bg());
    }
  }
  cout << ans;
  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
#define lli long long int
#define rep(i, n) for (lli i = 0; i < n; i++)
#define repb(i, n) for (lli i = n - 1; i >= 0; i--)
#define pb push_back
#define mp make_pair
#define bg() begin()
#define en() end()
#define ft first
#define sc second
#define vect_lli_it vector<lli>::iterator
#define set_lli_it set<lli>::iterator
#define inf 1e18
#define all(X) X.begin(), X.end()
#define iterate(X) for (it = X.begin(); it != X.end(); it++)
lli combi(lli n, lli k) {
  lli ans = 1;
  k = k > n - k ? n - k : k;
  lli j = 1;
  for (; j <= k; j++, n--) {
    if (n % j == 0) {
      ans *= n / j;
    } else if (ans % j == 0) {
      ans = ans / j * n;
    } else {
      ans = (ans * n) / j;
    }
  }
  return ans;
}
lli gcd(lli a, lli b) {
  if (b == 0)
    return a;
  return gcd(b, a % b);
}
lli findlcm(vector<lli> arr, lli n) {
  lli ans = arr[0];
  for (lli i = 1; i < n; i++)
    ans = (((arr[i] * ans)) / (gcd(arr[i], ans)));

  return ans;
}
lli power(lli x, lli y, lli p) {
  lli res = 1;
  x = x % p;
  if (x == 0)
    return 0;
  while (y > 0) {
    if (y & 1)
      res = (res * x) % p;
    y = y >> 1;
    x = (x * x) % p;
  }
  return res;
}
int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(NULL);
  lli n;
  cin >> n;
  vector<lli> arr;
  lli x;
  rep(i, n) {
    cin >> x;
    arr.pb(x);
  }
  vector<lli> sum_arr;
  rep(i, n) { sum_arr.pb(i + arr[i] + 1); }
  vector<lli> diff_arr;
  rep(i, n) { diff_arr.pb(i - arr[i] + 1); }
  sort(all(sum_arr));
  sort(all(diff_arr));
  vect_lli_it it;
  lli ans = 0;
  rep(i, sum_arr.size()) {
    it = lower_bound(all(diff_arr), sum_arr[i]);
    while (it != diff_arr.en() && *it == sum_arr[i]) {
      ans++;
      it++;
    }
  }
  cout << ans;
  return 0;
}