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

#include <bits/stdc++.h>
#define fo(i, a, b) for (int i = a; i < b; ++i)
#define rep(i, n) fo(i, 0, n)
#define pln(n) printf("%lld", n)
#define sll(n) scanf("%lld", &n)
#define ss(n) scanf("%s", n)
#define vi vector<int>
#define ll long long
#define pii pair<ll, ll>
#define pb push_back
#define mp make_pair
#define ull unsigned long long
#define INF 0x3f3f3f3f
#define dbg(x)                                                                 \\
  { cout << #x << ": " << (x) << endl; }
#define dbg2(x, y)                                                             \\
  { cout << #x << ": " << (x) << " , " << #y << ": " << (y) << endl; }
#define dbg3(x, y, z)                                                          \\
  {                                                                            \\
    cout << #x << ": " << (x) << " , " << #y << ": " << (y) << " , " << #z     \\
         << ": " << (z) << endl;                                               \\
  }
#define all(x) x.begin(), x.end()
#define sz(x) (int)x.size()
#define vll vector<long long>
#define vvll vector<vector<long long>>
#define vvvll vector<vector<vector<ll>>>
#define si set<int>
#define tr(c, it) for (decltype(c.begin()) it = c.begin(); it != c.end(); it++)
#define pis pair<int, string>
#define present(c, x) (c.find(x) != c.end())
#define cpresent(c, x) (find(all(c), x) != c.end())
#define ones(x) __builtin_popcount(x)
using namespace std;
int n;
struct compare {
  bool operator()(const pair<ll, ll> &value, const ll &key) {
    return (value.first < key);
  }
  bool operator()(const ll &key, const pair<ll, ll> &value) {
    return (key < value.first);
  }
};

// bool cmp(pii a, pii b){
//     if(a.second == b.second)return a.first < b.first;
//     return a.second > b.second;
// }

int main() {
  cin >> n;
  vll a(n);
  fo(i, 0, n) cin >> a[i];
  ll L[n];
  vector<pair<ll, ll>> R(n);
  fo(i, 0, n) {
    L[i] = i + 1 + a[i];
    R[i].first = i + 1 - a[i];
    R[i].second = i;
  }

  // fo(i,1,n+1)cout << L[i]<<" ";
  // cout << endl;
  // fo(i,1,n+1)cout <<"{"<<R[i].first<<", "<<R[i].second <<"}";
  sort(all(R));
  ll ans = 0;

  pair<std::vector<pii>::iterator, std::vector<pii>::iterator> p;
  for (int i = 0; i < n; i++) {
    p = equal_range(all(R), L[i], compare());
    for (auto it1 = p.first; it1 != p.second; it1++) {
      if (it1->second > i) {
        ans += p.second - it1;
        break;
      }
    }
  }
  cout << ans;
  // fo(i,0,n)cout <<"{"<<R[i].first<<", "<<R[i].second <<"}";
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
#define fo(i, a, b) for (int i = a; i < b; ++i)
#define rep(i, n) fo(i, 0, n)
#define pln(n) printf("%lld", n)
#define sll(n) scanf("%lld", &n)
#define ss(n) scanf("%s", n)
#define vi vector<int>
#define ll long long
#define pii pair<ll, ll>
#define pb push_back
#define mp make_pair
#define ull unsigned long long
#define INF 0x3f3f3f3f
#define dbg(x)                                                                 \\
  { cout << #x << ": " << (x) << endl; }
#define dbg2(x, y)                                                             \\
  { cout << #x << ": " << (x) << " , " << #y << ": " << (y) << endl; }
#define dbg3(x, y, z)                                                          \\
  {                                                                            \\
    cout << #x << ": " << (x) << " , " << #y << ": " << (y) << " , " << #z     \\
         << ": " << (z) << endl;                                               \\
  }
#define all(x) x.begin(), x.end()
#define sz(x) (int)x.size()
#define vll vector<long long>
#define vvll vector<vector<long long>>
#define vvvll vector<vector<vector<ll>>>
#define si set<int>
#define tr(c, it) for (decltype(c.begin()) it = c.begin(); it != c.end(); it++)
#define pis pair<int, string>
#define present(c, x) (c.find(x) != c.end())
#define cpresent(c, x) (find(all(c), x) != c.end())
#define ones(x) __builtin_popcount(x)
using namespace std;
int n;
struct compare {
  bool operator()(const pair<ll, ll> &value, const ll &key) {
    return (value.first < key);
  }
  bool operator()(const ll &key, const pair<ll, ll> &value) {
    return (key < value.first);
  }
};

// bool cmp(pii a, pii b){
//     if(a.second == b.second)return a.first < b.first;
//     return a.second > b.second;
// }

int main() {
  cin >> n;
  vll a(n);
  fo(i, 0, n) cin >> a[i];
  ll L[n];
  vector<pair<ll, ll>> R(n);
  fo(i, 0, n) {
    L[i] = i + 1 + a[i];
    R[i].first = i + 1 - a[i];
    R[i].second = i;
  }

  // fo(i,1,n+1)cout << L[i]<<" ";
  // cout << endl;
  // fo(i,1,n+1)cout <<"{"<<R[i].first<<", "<<R[i].second <<"}";
  sort(all(R));
  ll ans = 0;
  pair<std::vector<pii>::iterator, std::vector<pii>::iterator> p;
  for (int i = 0; i < n; i++) {
    p = equal_range(all(R), L[i], compare());
    for (auto it1 = p.first; it1 != p.second; it1++) {
      if (it1->second > i)
        ans++;
    }
  }
  cout << ans;
  // fo(i,0,n)cout <<"{"<<R[i].first<<", "<<R[i].second <<"}";
}