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

/*************************************************************************
    > File Name: solve.cpp
    > Author: liupo
    > Mail: lanzongwei@gmail.com
    > Created Time: 2020-05-03 21:01:27
 ************************************************************************/

#define GOODOJ
#define SYNC 0

#ifdef GOODOJ
#include <bits/stdc++.h>
#include <chrono>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/priority_queue.hpp>
#include <ext/pb_ds/tree_policy.hpp>
#include <ext/rope>
#include <random>
using namespace __gnu_pbds;
using namespace __gnu_cxx;
#else
#include <algorithm>
#include <cassert>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <deque>
#include <functional>
#include <iostream>
#include <iterator>
#include <limits>
#include <map>
#include <set>
#include <sstream>
#include <string>
#include <vector>
#endif
using namespace std;

struct InputOutputStream {
  enum { SIZE = 1000001 };
  char ibuf[SIZE], *s, *t, obuf[SIZE], *oh;
  bool eof;

  InputOutputStream() : s(), t(), oh(obuf), eof(false) {}
  ~InputOutputStream() { fwrite(obuf, 1, oh - obuf, stdout); }

  inline char read() {
    if (s == t)
      t = (s = ibuf) + fread(ibuf, 1, SIZE, stdin);
    return s == t ? -1 : *s++;
  }

  template <typename T> inline InputOutputStream &operator>>(T &x) {
    static char c;
    static bool iosig;
    for (c = read(), iosig = false; !isdigit(c); c = read()) {
      if (c == -1) {
        eof = true;
        return *this;
      }
      iosig |= c == '-';
    }
    for (x = 0; isdigit(c); c = read())
      x = x * 10 + (c ^ '0');
    if (iosig)
      x = -x;
    return *this;
  }

  inline void print(char c) {
    if (oh == obuf + SIZE) {
      fwrite(obuf, 1, SIZE, stdout);
      oh = obuf;
    }
    *oh++ = c;
  }

  template <typename T> inline void print(T x) {
    static int buf[23], cnt;
    if (x != 0) {
      if (x < 0)
        print('-'), x = -x;
      for (cnt = 0; x; x /= 10)
        buf[++cnt] = x % 10 | 48;
      while (cnt)
        print((char)buf[cnt--]);
    } else
      print('0');
  }

  template <typename T> inline InputOutputStream &operator<<(const T &x) {
    print(x);
    return *this;
  }

  inline void print(const char *x) {
    for (; *x; x++)
      print(*x);
  }
} io;

// #define cin io
// #define cout io

#define endl '\\n'
#define fep(i, b, e) for (int i = (b); i < (e); ++i)
#define rep(i, x) for (int i = 0; i < (x); ++i)
#define rap(i, x) for (auto &i : (x))
#define seg(t) (t).begin(), (t).end()
#define ep emplace_back
#define mkp make_pair
#define mkt make_tuple
#define qxx(i, x) for (int i = head[x]; ~i; i = node[i].nex)
#define F first
#define S second
#define lowbit(x) ((-(x)) & (x))
#define RE register
#define getchar() getchar_unlocked()
#ifdef DEBUG
void err(istream_iterator<string>) {}
template <typename T, typename... Args>
void err(istream_iterator<string> it, T a, Args... args) {
  cerr << *it << " = " << a << ' ';
  err(++it, args...);
}
#define debug(args...)                                                         \\
  {                                                                            \\
    string _s = #args;                                                         \\
    replace(seg(_s), ',', ' ');                                                \\
    cerr << "DEBUG:";                                                          \\
    istringstream it(_s);                                                      \\
    err(istream_iterator<string>(it), args);                                   \\
    cerr << endl;                                                              \\
  }
#else
#define debug(...)
#endif

template <typename T> inline bool cmax(T &a, const T &b) {
  return a < b ? a = b, 1 : 0;
}
template <typename T> inline bool cmin(T &a, const T &b) {
  return a > b ? a = b, 1 : 0;
}

#ifdef GOODOJ
mt19937 rng((uint32_t)chrono::steady_clock::now().time_since_epoch().count());
typedef __gnu_pbds::priority_queue<int> pq;
#endif
typedef std::string str;
typedef long long ll;
typedef double db;
typedef pair<int, int> pa;

const double P = acos(-1.0), eps = 1e-9;
struct point {
  db x, y;
};
inline int sign(db a) { return a < -eps ? -1 : a > eps; }
#define dot(p1, p2, p3)                                                        \\
  ((p2.x - p1.x) * (p3.x - p1.x) + (p2.y - p1.y) * (p3.y - p1.y))
#define cross(p1, p2, p3)                                                      \\
  ((p2.x - p1.x) * (p3.y - p1.y) - (p3.x - p1.x) * (p2.y - p1.y))
#define crossOp(p1, p2, p3) sign(cross(p1, p2, p3))

const int Ma = 2e5 + 100, inf = 0x3f3f3f3f, mod = 1e9 + 7;

int mz[Ma];

map<ll, ll> a, b;

signed main() {
#if SYNC == 1
  ios::sync_with_stdio(false);
  cin.tie(0);
#endif
  int n;
  cin >> n;
  rep(i, n) cin >> mz[i];
  rep(i, n) a[mz[i] + i]++, b[i - mz[i]]++;
  ll ans = 0;
  rap(i, a) if (b.count(i.F)) ans += i.S * b[i.F];

  cout << ans << endl;

  return 0;
}


Error: TLE

Fixed submission:

/*************************************************************************
    > File Name: solve.cpp
    > Author: liupo
    > Mail: lanzongwei@gmail.com
    > Created Time: 2020-05-03 21:01:27
 ************************************************************************/

#define GOODOJ
#define SYNC 0

#ifdef GOODOJ
#include <bits/stdc++.h>
#include <chrono>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/priority_queue.hpp>
#include <ext/pb_ds/tree_policy.hpp>
#include <ext/rope>
#include <random>
using namespace __gnu_pbds;
using namespace __gnu_cxx;
#else
#include <algorithm>
#include <cassert>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <deque>
#include <functional>
#include <iostream>
#include <iterator>
#include <limits>
#include <map>
#include <set>
#include <sstream>
#include <string>
#include <vector>
#endif
using namespace std;

struct InputOutputStream {
  enum { SIZE = 1000001 };
  char ibuf[SIZE], *s, *t, obuf[SIZE], *oh;
  bool eof;

  InputOutputStream() : s(), t(), oh(obuf), eof(false) {}
  ~InputOutputStream() { fwrite(obuf, 1, oh - obuf, stdout); }

  inline char read() {
    if (s == t)
      t = (s = ibuf) + fread(ibuf, 1, SIZE, stdin);
    return s == t ? -1 : *s++;
  }

  template <typename T> inline InputOutputStream &operator>>(T &x) {
    static char c;
    static bool iosig;
    for (c = read(), iosig = false; !isdigit(c); c = read()) {
      if (c == -1) {
        eof = true;
        return *this;
      }
      iosig |= c == '-';
    }
    for (x = 0; isdigit(c); c = read())
      x = x * 10 + (c ^ '0');
    if (iosig)
      x = -x;
    return *this;
  }

  inline void print(char c) {
    if (oh == obuf + SIZE) {
      fwrite(obuf, 1, SIZE, stdout);
      oh = obuf;
    }
    *oh++ = c;
  }

  template <typename T> inline void print(T x) {
    static int buf[23], cnt;
    if (x != 0) {
      if (x < 0)
        print('-'), x = -x;
      for (cnt = 0; x; x /= 10)
        buf[++cnt] = x % 10 | 48;
      while (cnt)
        print((char)buf[cnt--]);
    } else
      print('0');
  }

  template <typename T> inline InputOutputStream &operator<<(const T &x) {
    print(x);
    return *this;
  }

  inline void print(const char *x) {
    for (; *x; x++)
      print(*x);
  }
} io;

#define cin io
#define cout io

#define endl '\\n'
#define fep(i, b, e) for (int i = (b); i < (e); ++i)
#define rep(i, x) for (int i = 0; i < (x); ++i)
#define rap(i, x) for (auto &i : (x))
#define seg(t) (t).begin(), (t).end()
#define ep emplace_back
#define mkp make_pair
#define mkt make_tuple
#define qxx(i, x) for (int i = head[x]; ~i; i = node[i].nex)
#define F first
#define S second
#define lowbit(x) ((-(x)) & (x))
#define RE register
#define getchar() getchar_unlocked()
#ifdef DEBUG
void err(istream_iterator<string>) {}
template <typename T, typename... Args>
void err(istream_iterator<string> it, T a, Args... args) {
  cerr << *it << " = " << a << ' ';
  err(++it, args...);
}
#define debug(args...)                                                         \\
  {                                                                            \\
    string _s = #args;                                                         \\
    replace(seg(_s), ',', ' ');                                                \\
    cerr << "DEBUG:";                                                          \\
    istringstream it(_s);                                                      \\
    err(istream_iterator<string>(it), args);                                   \\
    cerr << endl;                                                              \\
  }
#else
#define debug(...)
#endif

template <typename T> inline bool cmax(T &a, const T &b) {
  return a < b ? a = b, 1 : 0;
}
template <typename T> inline bool cmin(T &a, const T &b) {
  return a > b ? a = b, 1 : 0;
}

#ifdef GOODOJ
mt19937 rng((uint32_t)chrono::steady_clock::now().time_since_epoch().count());
typedef __gnu_pbds::priority_queue<int> pq;
#endif
typedef std::string str;
typedef long long ll;
typedef double db;
typedef pair<int, int> pa;

const double P = acos(-1.0), eps = 1e-9;
struct point {
  db x, y;
};
inline int sign(db a) { return a < -eps ? -1 : a > eps; }
#define dot(p1, p2, p3)                                                        \\
  ((p2.x - p1.x) * (p3.x - p1.x) + (p2.y - p1.y) * (p3.y - p1.y))
#define cross(p1, p2, p3)                                                      \\
  ((p2.x - p1.x) * (p3.y - p1.y) - (p3.x - p1.x) * (p2.y - p1.y))
#define crossOp(p1, p2, p3) sign(cross(p1, p2, p3))

const int Ma = 2e5 + 100, inf = 0x3f3f3f3f, mod = 1e9 + 7;

int mz[Ma], id[Ma];

bool cmp(int a, int b) { return mz[a] < mz[b]; }

signed main() {
#if SYNC == 1
  ios::sync_with_stdio(false);
  cin.tie(0);
#endif
  int n;
  cin >> n;
  rep(i, n) cin >> mz[i], id[i] = i;
  int ans = 0;
  sort(id, id + n, cmp);
  rep(i, n) {
    for (int j = i + mz[i]; j < n; j++)
      if (mz[i] + mz[j] == j - i)
        ++ans;
  }
  cout << ans << endl;

  return 0;
}