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

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

typedef long long int ll;
typedef long double ld;
typedef double db;
typedef string str;
typedef bool bl;

typedef pair<int, int> pi;
typedef pair<ll, ll> pl;
typedef pair<ld, ld> pd;

typedef vector<int> vi;
typedef vector<ll> vl;
typedef vector<ld> vd;
typedef vector<str> vs;
typedef vector<pi> vpi;
typedef vector<pl> vpl;
typedef vector<pd> vpd;
typedef vector<vector<ll>> vvl;
typedef vector<bl> vbl;

#define mp make_pair
#define ff first
#define ss second
#define sz(x) (int)x.size()
#define all(x) begin(x), end(x)
#define rall(x) (x).rbegin(), (x).rend()
#define rsz resize
/// #define ins insert
#define ft front()
#define bk back()
#define pf push_front
#define pb push_back
#define eb emplace_back
#define lb lower_bound
#define ub upper_bound

#define FOR(i, a, b) for (ll i = (a); i < (b); ++i)
#define F0R(i, a) FOR(i, 0, a)
#define ROF(i, a, b) for (ll i = (b)-1; i >= (a); --i)
#define R0F(i, a) ROF(i, 0, a)
#define trav(a, x) for (auto &a : x)

const ll mod = 1e6 + 3; // 998244353;
const int MX = 2e5 + 5;
const ll inf = 1e18;
const ld PI = acos((ld)-1);
const int xd[4] = {1, 0, -1, 0}, yd[4] = {0, 1, 0, -1};

template <class T> bool ckmin(T &a, const T &b) { return b < a ? a = b, 1 : 0; }
template <class T> bool ckmax(T &a, const T &b) { return a < b ? a = b, 1 : 0; }
int pct(int x) { return __builtin_popcount(x); }

namespace input {
template <class T> void re(complex<T> &x);
template <class T1, class T2> void re(pair<T1, T2> &p);
template <class T> void re(vector<T> &a);
template <class T, size_t SZ> void re(array<T, SZ> &a);

template <class T> void re(T &x) { cin >> x; }
void re(double &x) {
  string t;
  re(t);
  x = stod(t);
}
void re(ld &x) {
  string t;
  re(t);
  x = stold(t);
}
template <class T, class... Ts> void re(T &t, Ts &...ts) {
  re(t);
  re(ts...);
}

template <class T> void re(complex<T> &x) {
  T a, b;
  re(a, b);
  x = {a, b};
}
template <class T1, class T2> void re(pair<T1, T2> &p) { re(p.f, p.s); }
template <class T> void re(vector<T> &a) { F0R(i, sz(a)) re(a[i]); }
template <class T, size_t SZ> void re(array<T, SZ> &a) { F0R(i, SZ) re(a[i]); }
} // namespace input

using namespace input;

namespace output {
void pr(int x) { cout << x; }
void pr(long x) { cout << x; }
void pr(ll x) { cout << x; }
void pr(unsigned x) { cout << x; }
void pr(unsigned long x) { cout << x; }
void pr(unsigned long long x) { cout << x; }
void pr(float x) { cout << x; }
void pr(double x) { cout << x; }
void pr(ld x) { cout << x; }
void pr(char x) { cout << x; }
void pr(const char *x) { cout << x; }
void pr(const string &x) { cout << x; }
void pr(bool x) { pr(x ? "true" : "false"); }
template <class T> void pr(const complex<T> &x) { cout << x; }

template <class T1, class T2> void pr(const pair<T1, T2> &x);
template <class T> void pr(const T &x);

template <class T, class... Ts> void pr(const T &t, const Ts &...ts) {
  pr(t);
  pr(ts...);
}
template <class T1, class T2> void pr(const pair<T1, T2> &x) {
  pr("{", x.f, ", ", x.s, "}");
}
template <class T> void pr(const T &x) {
  pr("{"); // const iterator needed for vector<bool>
  bool fst = 1;
  for (const auto &a : x)
    pr(!fst ? ", " : "", a), fst = 0;
  pr("}");
}

void ps() { pr("\\n"); } // print w/ spaces
template <class T, class... Ts> void ps(const T &t, const Ts &...ts) {
  pr(t);
  if (sizeof...(ts))
    pr(" ");
  ps(ts...);
}

void pc() { pr("]\\n"); } // debug w/ commas
template <class T, class... Ts> void pc(const T &t, const Ts &...ts) {
  pr(t);
  if (sizeof...(ts))
    pr(", ");
  pc(ts...);
}
#define dbg(x...) pr("[", #x, "] = ["), pc(x);
} // namespace output

using namespace output;

namespace io {
void setIn(string s) { freopen(s.c_str(), "r", stdin); }
void setOut(string s) { freopen(s.c_str(), "w", stdout); }
void setIO(string s = "") {
  ios_base::sync_with_stdio(0);
  cin.tie(0); // fast I/O
  // cin.exceptions(cin.failbit);
  // throws exception when do smth illegal
  // ex. try to read letter into int
  if (sz(s)) {
    setIn(s + ".in"), setOut(s + ".out");
  } // for USACO
}
} // namespace io

using namespace io;

mt19937 rng((uint32_t)chrono::steady_clock::now().time_since_epoch().count());

inline ll mul(ll a, ll b) { return ((a % mod) * (b % mod)) % mod; }
inline ll add(ll a, ll b) {
  a += mod;
  b += mod;
  return (a % mod + b % mod) % mod;
}
ll binpow(ll x, ll y) {
  ll z = 1;
  while (y) {
    if (y & 1)
      z = mul(z, x);
    x = mul(x, x);
    y >>= 1;
  }
  return z;
}

inline ll inv(ll x) { return binpow(x, mod - 2); }
inline ll divide(ll x, ll y) { return mul(x, inv(y)); }
inline ll lcm(ll a, ll b) { return (a * b) / (__gcd(a, b)); }
inline ll gcd(ll a, ll b) { return __gcd(a, b); }
/*
unique(all(vec)); will remove all duplicate elements from the vector
find_by_order() - Returns an iterator to the k-th largest element (counting from
zero) order_of_key()  - The number of items in a set that are strictly smaller
than our item *max_element (first_iterator, last_iterator) – To find the maximum
element of a vector. *min_element (first_iterator, last_iterator) – To find the
minimum element of a vector. accumulate(first_iterator, last_iterator, initial
value of sum) – Does the summation of vector elements count(first_iterator,
last_iterator,x) – To count the occurrences of x in vector. find(first_iterator,
last_iterator, x) – Points to last address of vector ((name_of_vector).end()) if
element is not present in vector. binary_search(first_iterator, last_iterator,
x) – Tests whether x exists in sorted vector or not. lower_bound(first_iterator,
last_iterator, x) – returns an iterator pointing to the first element in the
range [first,last) which has a value not less than ‘x’.
upper_bound(first_iterator, last_iterator, x) – returns an iterator pointing to
the first element in the range [first,last) which has a value greater than ‘x’.
arr.erase(position to be deleted) – This erases selected element in vector and
shifts and resizes the vector elements accordingly.
arr.erase(unique(arr.begin(),arr.end()),arr.end()) – This erases the duplicate
occurrences in sorted vector in a single line. next_permutation(first_iterator,
last_iterator) – This modified the vector to its next permutation.
prev_permutation(first_iterator, last_iterator) – This modified the vector to
its previous permutation.
 */

int main() {
  ll n;
  re(n);
  vl a(n);
  re(a);
  vpl pp;
  for (ll i = 0; i < n; i++) {
    if (a[i] < n - 1) {
      pp.pb({i + 1, a[i]});
    }
  }

  ll p[3 * n + 1000] = {0};

  vl m(n, 0);
  for (ll i = 0; i < pp.size(); i++) {
    auto it = pp[i];
    ll id = it.first;
    ll val = it.second;
    m[i] = id - val;
    p[id + val]++;
  }
  ll ans = 0;
  for (ll i = 0; i < pp.size(); i++) {
    if (m[i] > 0)
      ans += p[m[i]];
  }

  cout << ans;
}

Error: TLE

Fixed submission:

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

typedef long long int ll;
typedef long double ld;
typedef double db;
typedef string str;
typedef bool bl;

typedef pair<int, int> pi;
typedef pair<ll, ll> pl;
typedef pair<ld, ld> pd;

typedef vector<int> vi;
typedef vector<ll> vl;
typedef vector<ld> vd;
typedef vector<str> vs;
typedef vector<pi> vpi;
typedef vector<pl> vpl;
typedef vector<pd> vpd;
typedef vector<vector<ll>> vvl;
typedef vector<bl> vbl;

#define mp make_pair
#define ff first
#define ss second
#define sz(x) (int)x.size()
#define all(x) begin(x), end(x)
#define rall(x) (x).rbegin(), (x).rend()
#define rsz resize
/// #define ins insert
#define ft front()
#define bk back()
#define pf push_front
#define pb push_back
#define eb emplace_back
#define lb lower_bound
#define ub upper_bound

#define FOR(i, a, b) for (ll i = (a); i < (b); ++i)
#define F0R(i, a) FOR(i, 0, a)
#define ROF(i, a, b) for (ll i = (b)-1; i >= (a); --i)
#define R0F(i, a) ROF(i, 0, a)
#define trav(a, x) for (auto &a : x)

const ll mod = 1e6 + 3; // 998244353;
const int MX = 2e5 + 5;
const ll inf = 1e18;
const ld PI = acos((ld)-1);
const int xd[4] = {1, 0, -1, 0}, yd[4] = {0, 1, 0, -1};

template <class T> bool ckmin(T &a, const T &b) { return b < a ? a = b, 1 : 0; }
template <class T> bool ckmax(T &a, const T &b) { return a < b ? a = b, 1 : 0; }
int pct(int x) { return __builtin_popcount(x); }

namespace input {
template <class T> void re(complex<T> &x);
template <class T1, class T2> void re(pair<T1, T2> &p);
template <class T> void re(vector<T> &a);
template <class T, size_t SZ> void re(array<T, SZ> &a);

template <class T> void re(T &x) { cin >> x; }
void re(double &x) {
  string t;
  re(t);
  x = stod(t);
}
void re(ld &x) {
  string t;
  re(t);
  x = stold(t);
}
template <class T, class... Ts> void re(T &t, Ts &...ts) {
  re(t);
  re(ts...);
}

template <class T> void re(complex<T> &x) {
  T a, b;
  re(a, b);
  x = {a, b};
}
template <class T1, class T2> void re(pair<T1, T2> &p) { re(p.f, p.s); }
template <class T> void re(vector<T> &a) { F0R(i, sz(a)) re(a[i]); }
template <class T, size_t SZ> void re(array<T, SZ> &a) { F0R(i, SZ) re(a[i]); }
} // namespace input

using namespace input;

namespace output {
void pr(int x) { cout << x; }
void pr(long x) { cout << x; }
void pr(ll x) { cout << x; }
void pr(unsigned x) { cout << x; }
void pr(unsigned long x) { cout << x; }
void pr(unsigned long long x) { cout << x; }
void pr(float x) { cout << x; }
void pr(double x) { cout << x; }
void pr(ld x) { cout << x; }
void pr(char x) { cout << x; }
void pr(const char *x) { cout << x; }
void pr(const string &x) { cout << x; }
void pr(bool x) { pr(x ? "true" : "false"); }
template <class T> void pr(const complex<T> &x) { cout << x; }

template <class T1, class T2> void pr(const pair<T1, T2> &x);
template <class T> void pr(const T &x);

template <class T, class... Ts> void pr(const T &t, const Ts &...ts) {
  pr(t);
  pr(ts...);
}
template <class T1, class T2> void pr(const pair<T1, T2> &x) {
  pr("{", x.f, ", ", x.s, "}");
}
template <class T> void pr(const T &x) {
  pr("{"); // const iterator needed for vector<bool>
  bool fst = 1;
  for (const auto &a : x)
    pr(!fst ? ", " : "", a), fst = 0;
  pr("}");
}

void ps() { pr("\\n"); } // print w/ spaces
template <class T, class... Ts> void ps(const T &t, const Ts &...ts) {
  pr(t);
  if (sizeof...(ts))
    pr(" ");
  ps(ts...);
}

void pc() { pr("]\\n"); } // debug w/ commas
template <class T, class... Ts> void pc(const T &t, const Ts &...ts) {
  pr(t);
  if (sizeof...(ts))
    pr(", ");
  pc(ts...);
}
#define dbg(x...) pr("[", #x, "] = ["), pc(x);
} // namespace output

using namespace output;

namespace io {
void setIn(string s) { freopen(s.c_str(), "r", stdin); }
void setOut(string s) { freopen(s.c_str(), "w", stdout); }
void setIO(string s = "") {
  ios_base::sync_with_stdio(0);
  cin.tie(0); // fast I/O
  // cin.exceptions(cin.failbit);
  // throws exception when do smth illegal
  // ex. try to read letter into int
  if (sz(s)) {
    setIn(s + ".in"), setOut(s + ".out");
  } // for USACO
}
} // namespace io

using namespace io;

mt19937 rng((uint32_t)chrono::steady_clock::now().time_since_epoch().count());

inline ll mul(ll a, ll b) { return ((a % mod) * (b % mod)) % mod; }
inline ll add(ll a, ll b) {
  a += mod;
  b += mod;
  return (a % mod + b % mod) % mod;
}
ll binpow(ll x, ll y) {
  ll z = 1;
  while (y) {
    if (y & 1)
      z = mul(z, x);
    x = mul(x, x);
    y >>= 1;
  }
  return z;
}

inline ll inv(ll x) { return binpow(x, mod - 2); }
inline ll divide(ll x, ll y) { return mul(x, inv(y)); }
inline ll lcm(ll a, ll b) { return (a * b) / (__gcd(a, b)); }
inline ll gcd(ll a, ll b) { return __gcd(a, b); }
/*
unique(all(vec)); will remove all duplicate elements from the vector
find_by_order() - Returns an iterator to the k-th largest element (counting from
zero) order_of_key()  - The number of items in a set that are strictly smaller
than our item *max_element (first_iterator, last_iterator) – To find the maximum
element of a vector. *min_element (first_iterator, last_iterator) – To find the
minimum element of a vector. accumulate(first_iterator, last_iterator, initial
value of sum) – Does the summation of vector elements count(first_iterator,
last_iterator,x) – To count the occurrences of x in vector. find(first_iterator,
last_iterator, x) – Points to last address of vector ((name_of_vector).end()) if
element is not present in vector. binary_search(first_iterator, last_iterator,
x) – Tests whether x exists in sorted vector or not. lower_bound(first_iterator,
last_iterator, x) – returns an iterator pointing to the first element in the
range [first,last) which has a value not less than ‘x’.
upper_bound(first_iterator, last_iterator, x) – returns an iterator pointing to
the first element in the range [first,last) which has a value greater than ‘x’.
arr.erase(position to be deleted) – This erases selected element in vector and
shifts and resizes the vector elements accordingly.
arr.erase(unique(arr.begin(),arr.end()),arr.end()) – This erases the duplicate
occurrences in sorted vector in a single line. next_permutation(first_iterator,
last_iterator) – This modified the vector to its next permutation.
prev_permutation(first_iterator, last_iterator) – This modified the vector to
its previous permutation.
 */

int main() {
  ll n;
  re(n);
  vl a(n);
  re(a);
  vl p(n, 0);
  vl m(n, 0);
  for (ll i = 1; i <= n; i++) {
    p[i - 1] = i + a[i - 1];
    m[i - 1] = i - a[i - 1];
  }
  ll ans = 0;
  sort(all(p));
  trav(t, m) { ans += count(all(p), t); }
  cout << ans;
}