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

#include "bits/stdc++.h"

#define REP(i, num) for (int i = 0; i < (num); ++i)
#define LOOP(i) while (i--)
#define ALL(c) c.begin(), c.end()
#define PRINTALL(c)                                                            \\
  for (auto pitr = c.begin(); pitr != c.end(); ++pitr) {                       \\
    cout << *pitr;                                                             \\
    if (next(pitr, 1) != c.end())                                              \\
      cout << ' ';                                                             \\
  }                                                                            \\
  cout << endl;
#define PAIRCOMP(c, comp)                                                      \\
  [](const pair<ll, ll> &lhs, const pair<ll, ll> &rhs) {                       \\
    return lhs.c comp rhs.c;                                                   \\
  }

using namespace std;
using ll = long long;

constexpr ll atcoder_mod = 1e9 + 7;

template <typename T = int> T in() {
  T x;
  cin >> x;
  return (x);
}
template <typename T = int, typename C = vector<T>> C vecin(int N) {
  C x(N);
  REP(i, N) { x[i] = in<T>(); }
  return move(x);
}

void vout() { cout << endl; }
template <typename Head, typename... Tail> void vout(Head &&h, Tail &&...t) {
  cout << ' ' << h;
  vout(forward<Tail>(t)...);
}
void out() { cout << endl; }
template <typename Head, typename... Tail> void out(Head &&h, Tail &&...t) {
  cout << h;
  vout(forward<Tail>(t)...);
}

int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);
  cout << fixed << setprecision(10);

  int N = in();
  vector<ll> A = vecin<ll>(N);
  vector<ll> B(N), C(N);

  map<ll, ll> MB, MC;
  REP(i, N) {
    if (A[i] >= N)
      continue;
    B[i] = i - A[i];
    C[i] = A[i] + i;
    MB[B[i]]++;
    MC[C[i]]++;
  }

  ll sum = 0;
  for (auto &x : MB) {
    ll V = x.first;
    sum += x.second * MC[V];
  }
  out(sum);
  return 0;
}


Error: TLE

Fixed submission:

#include "bits/stdc++.h"

#define REP(i, num) for (int i = 0; i < (num); ++i)
#define LOOP(i) while (i--)
#define ALL(c) c.begin(), c.end()
#define PRINTALL(c)                                                            \\
  for (auto pitr = c.begin(); pitr != c.end(); ++pitr) {                       \\
    cout << *pitr;                                                             \\
    if (next(pitr, 1) != c.end())                                              \\
      cout << ' ';                                                             \\
  }                                                                            \\
  cout << endl;
#define PAIRCOMP(c, comp)                                                      \\
  [](const pair<ll, ll> &lhs, const pair<ll, ll> &rhs) {                       \\
    return lhs.c comp rhs.c;                                                   \\
  }

using namespace std;
using ll = long long;

constexpr ll atcoder_mod = 1e9 + 7;

template <typename T = int> T in() {
  T x;
  cin >> x;
  return (x);
}
template <typename T = int, typename C = vector<T>> C vecin(int N) {
  C x(N);
  REP(i, N) { x[i] = in<T>(); }
  return move(x);
}

void vout() { cout << endl; }
template <typename Head, typename... Tail> void vout(Head &&h, Tail &&...t) {
  cout << ' ' << h;
  vout(forward<Tail>(t)...);
}
void out() { cout << endl; }
template <typename Head, typename... Tail> void out(Head &&h, Tail &&...t) {
  cout << h;
  vout(forward<Tail>(t)...);
}

int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);
  cout << fixed << setprecision(10);

  int N = in();
  vector<ll> A = vecin<ll>(N);

  vector<vector<ll>> C(N);
  REP(i, N) {
    if (A[i] < N)
      C[A[i]].push_back(i);
  }

  int sum = 0;
  REP(i, N) {
    if (A[i] >= N)
      continue;
    for (int j = i + 1 + A[i], k = 1; j < N; j++, k++) {
      if (binary_search(ALL(C[k]), j))
        sum++;
    }
  }

  out(sum);
  return 0;
}