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

#include <bits/stdc++.h>
using namespace std;

using ll = long long;
using ul = unsigned long;
using ull = unsigned long long;
#define vec vector
#define rep(_i, _n) for (ll _i = 0; _i < _n; _i++)
#define rep2(_i, _s, _n) for (ll _i = _s; _i < _n; _i++)
#define elif else if
#define push(_a) push_back(_a)
#define append(_x) push_back(_x)
#define fullSort(_a) sort(_a.begin(), _a.end())
#define inputVecI(_name, _n)                                                   \\
  vec<int> _name(_n);                                                          \\
  rep(_i, _n) cin >> _name[_i];
#define inputVecL(_name, _n)                                                   \\
  vec<long> _name(_n);                                                         \\
  rep(_i, _n) cin >> _name[_i];
#define inputVecLL(_name, _n)                                                  \\
  vec<ll> _name(_n);                                                           \\
  rep(_i, _n) cin >> _name[_i];
#define inputVecUL(_name, _n)                                                  \\
  vec<ul> _name(_n);                                                           \\
  rep(_i, _n) cin >> _name[_i];
#define inputVecULL(_name, _n)                                                 \\
  vec<ull> _name(_n);                                                          \\
  rep(_i, _n) cin >> _name[_i];
#define inputVecS(_name, _n)                                                   \\
  vec<string> _name(_n);                                                       \\
  rep(_i, _n) cin >> _name[_i];
#define findMaxIdx(_name)                                                      \\
  distance(_name.begin(), max_element(_name.begin(), _name.end()))
#define printVec(_vec)                                                         \\
  cout << "{";                                                                 \\
  rep(i, _vec.size() - 1) cout << _vec[i] << ",";                              \\
  cout << _vec[_vec.size() - 1];                                               \\
  cout << "}\\n";
#define mapLoop(_i, _name) for (auto i = _name.begin(); i != _name.end(); i++)

#define comp_idx(_name)                                                        \\
  [&_name](long l_idx, long r_idx) { return _name[l_idx] < _name[r_idx]; }

int main() {
  long N;
  cin >> N;
  inputVecLL(A, N);
  map<ll, long> ansbuf;
  vec<ll> SA(N);
  rep(i, N) {
    ansbuf[(i + 1) + A[i]] += 1;
    SA[i] = (i + 1) - A[i];
  }

  long ans = 0;
  rep(i, N) { ans += ansbuf[SA[i]]; }

  cout << ans << endl;
  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;

using ll = long long;
using ul = unsigned long;
using ull = unsigned long long;
#define vec vector
#define rep(_i, _n) for (ll _i = 0; _i < _n; _i++)
#define rep2(_i, _s, _n) for (ll _i = _s; _i < _n; _i++)
#define elif else if
#define push(_a) push_back(_a)
#define append(_x) push_back(_x)
#define fullSort(_a) sort(_a.begin(), _a.end())
#define inputVecI(_name, _n)                                                   \\
  vec<int> _name(_n);                                                          \\
  rep(_i, _n) cin >> _name[_i];
#define inputVecL(_name, _n)                                                   \\
  vec<long> _name(_n);                                                         \\
  rep(_i, _n) cin >> _name[_i];
#define inputVecLL(_name, _n)                                                  \\
  vec<ll> _name(_n);                                                           \\
  rep(_i, _n) cin >> _name[_i];
#define inputVecUL(_name, _n)                                                  \\
  vec<ul> _name(_n);                                                           \\
  rep(_i, _n) cin >> _name[_i];
#define inputVecULL(_name, _n)                                                 \\
  vec<ull> _name(_n);                                                          \\
  rep(_i, _n) cin >> _name[_i];
#define inputVecS(_name, _n)                                                   \\
  vec<string> _name(_n);                                                       \\
  rep(_i, _n) cin >> _name[_i];
#define findMaxIdx(_name)                                                      \\
  distance(_name.begin(), max_element(_name.begin(), _name.end()))
#define printVec(_vec)                                                         \\
  cout << "{";                                                                 \\
  rep(i, _vec.size() - 1) cout << _vec[i] << ",";                              \\
  cout << _vec[_vec.size() - 1];                                               \\
  cout << "}\\n";
#define mapLoop(_i, _name) for (auto i = _name.begin(); i != _name.end(); i++)

#define comp_idx(_name)                                                        \\
  [&_name](long l_idx, long r_idx) { return _name[l_idx] < _name[r_idx]; }

int main() {
  long N;
  cin >> N;
  inputVecLL(A, N);
  vec<ll> WA0(N);
  vec<ll> WA(N);
  vec<long> idxWA(N);
  vec<ll> SA0(N);
  vec<ll> SA(N);
  vec<long> idxSA(N);
  rep(i, N) {
    WA0[i] = (i + 1) + A[i];
    idxWA[i] = i;
    SA0[i] = (i + 1) - A[i];
    idxSA[i] = i;
  }

  sort(idxWA.begin(), idxWA.end(), comp_idx(WA0));
  sort(idxSA.begin(), idxSA.end(), comp_idx(SA0));
  rep(i, N) {
    WA[i] = WA0[idxWA[i]];
    SA[i] = SA0[idxSA[i]];
  }

  long ans = 0;
  long bufidx = 0;
  rep(i, N) {
    ll buf = WA[i];
    bufidx =
        distance(SA.begin(), lower_bound(SA.begin() + bufidx, SA.end(), buf));
    rep2(j, bufidx, N) {
      if (buf == SA[j]) {
        ans++;
      } else {
        break;
      }
    }
  }

  cout << ans << endl;
  return 0;
}