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

#include <algorithm>
#include <bitset>
#include <cctype>
#include <climits>
#include <cmath>
#include <cstring>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <type_traits>
#include <vector>

#include <list>
using namespace std;

#define REP(i, n) for (long long i = 0; i < (n); ++i)
#define rep(i, a, b) for (long long i = a; i < (b); ++i)
#define YES(j) cout << (j ? "YES" : "NO") << endl;
#define Yes(j) std::cout << (j ? "Yes" : "No") << endl;

typedef long long ll;

int main(void) {
  map<int, vector<int>> mp;
  map<int, vector<int>> mp2;
  int n;
  cin >> n;
  int a[n];
  REP(i, n)
  cin >> a[i];
  long long ans = 0;
  REP(i, n) {
    // i + 1 + Ai
    int x = i + 1 + a[i];
    mp[x].push_back(i + 1);
    // i + 1 - Ai
    x = i + 1 - a[i];
    mp2[x].push_back(i + 1);
  }

  for (auto i : mp) {
    sort(i.second.begin(), i.second.end());
  }
  for (auto i : mp2) {
    sort(i.second.begin(), i.second.end());
  }

  for (auto i : mp) {
    int x = i.first;
    // cout << "x " << x << endl;
    vector<int> pair = i.second;
    for (auto j : pair) {
      int val = j;
      // cout << "i " << val << endl;

      auto it = upper_bound(mp2[x].begin(), mp2[x].end(), val);
      int dist = distance(it, mp2[x].end());
      ans += dist;
      /*
      for (int k = 0; k < mp2[x].size(); k++)
      {
          if (mp2[x][k] > val)
          {
              //cout << "j " << k << endl;
              ans++;
          }
      }
      */
    }
  }

  cout << ans << endl;
}

Error: TLE

Fixed submission:

#include <algorithm>
#include <bitset>
#include <cctype>
#include <climits>
#include <cmath>
#include <cstring>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <type_traits>
#include <vector>

#include <list>
using namespace std;

#define REP(i, n) for (long long i = 0; i < (n); ++i)
#define rep(i, a, b) for (long long i = a; i < (b); ++i)
#define YES(j) cout << (j ? "YES" : "NO") << endl;
#define Yes(j) std::cout << (j ? "Yes" : "No") << endl;

typedef long long ll;

int main(void) {
  map<int, vector<int>> mp;
  map<int, vector<int>> mp2;
  int n;
  cin >> n;
  int a[n];
  REP(i, n)
  cin >> a[i];
  long long ans = 0;
  REP(i, n) {
    // i + 1 + Ai
    int x = i + 1 + a[i];
    mp[x].push_back(i + 1);
    // i + 1 - Ai
    x = i + 1 - a[i];
    mp2[x].push_back(i + 1);
  }

  for (auto i : mp) {
    sort(i.second.begin(), i.second.end());
  }
  for (auto i : mp2) {
    sort(i.second.begin(), i.second.end());
  }

  for (auto i : mp) {
    int x = i.first;
    // cout << "x " << x << endl;
    vector<int> pair = i.second;
    for (auto j : pair) {
      int val = j;
      // cout << "i " << val << endl;

      /*
      auto it = upper_bound(mp2[x].begin(), mp2[x].end(), val);
      if (it != mp2[x].end())
      {
          int dist = distance(mp2[x].begin(), it);
          ans += dist;
      }
      */
      for (int k = 0; k < mp2[x].size(); k++) {
        if (mp2[x][k] > val) {
          // cout << "j " << k << endl;
          ans++;
        }
      }
    }
  }

  cout << ans << endl;
}