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

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

#define pii pair<int, int>
#define vi vector<int>
#define pb push_back
#define eb emplace_back
#define mp make_pair
#define x first
#define y second
#define sz(x) (x).size()
#define all(x) (x).begin(), (x).end()  // Forward traversal
#define rall(x) (x).rbegin, (x).rend() // reverse traversal

#define ll long long
#define ld long double

const ld EPS = 10e-9;
const int INF = 1e9 + 5;
const long long INFLL = (ll)INF * (ll)INF;
const ll MOD = 1e9 + 7;
const ld PI = 3.14159265358979323846;

template <typename T> ostream &operator<<(ostream &os, const vector<T> &v) {
  os << "[";
  for (int i = 0; i < (int)v.size(); ++i) {
    os << v[i];
    if (i != (int)v.size() - 1)
      os << ", ";
  }
  os << "]\\n";
  return os;
}

template <typename T> ostream &operator<<(ostream &os, const set<T> &v) {
  os << "{";
  for (const auto &it : v) {
    os << it;
    if (it != *v.rbegin())
      os << ", ";
  }
  os << "}\\n";
  return os;
}

template <typename T>
ostream &operator<<(ostream &os, const unordered_set<T> &v) {
  os << "{";
  for (const auto &it : v) {
    os << it;
    if (it != *v.rbegin())
      os << ", ";
  }
  os << "}\\n";
  return os;
}

template <typename T, typename S>
ostream &operator<<(ostream &os, const map<T, S> &v) {
  os << "{";
  for (const auto &it : v) {
    os << it.first << " : " << it.second;

    if (it != *v.rbegin())
      os << ", ";
  }

  os << "}\\n";
  return os;
}

template <typename T, typename S>
ostream &operator<<(ostream &os, const unordered_map<T, S> &v) {
  os << "{";
  for (const auto &it : v) {
    os << it.first << " : " << it.second;

    if (it != *v.rbegin())
      os << ", ";
  }

  os << "}\\n";
  return os;
}

template <typename T, typename S>
ostream &operator<<(ostream &os, const pair<T, S> &v) {
  os << "(";
  os << v.first << ", " << v.second << ")";
  os << "\\n";
  return os;
}

#define DEBUG

#ifdef DEBUG
#define debug(...) __f(#__VA_ARGS__, __VA_ARGS__)
template <typename Arg1> void __f(const char *name, Arg1 &&arg1) {
  cerr << name << " : " << arg1;
}
template <typename Arg1, typename... Args>
void __f(const char *names, Arg1 &&arg1, Args &&...args) {
  const char *comma = strchr(names + 1, ',');
  cerr.write(names, comma - names) << " : " << arg1 << " | ";
  __f(comma + 1, args...);
}
#else
#define debug(...)
#endif

// __gcd(x, y)

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(0);

  int N;

  cin >> N;

  vector<ll> a(N + 1);
  for (int i = 1; i <= N; i++)
    cin >> a[i];

  unordered_map<ll, ll> s1, s2;

  ll ma = 0;
  for (int i = 1; i <= N; i++) {
    ma = max(ma, i + a[i]);
    s1[i + a[i]]++;
    s2[i - a[i]]++;
  }

  ll num = 0;
  for (int i = 1; i <= N; i++) {
    if (s1.count(i) && s2.count(i))
      num += s1[i] * s2[i];
  }

  cout << num << endl;

  //    vector<int> indices;
  //    for(int i = 0; i <= N; i++)indices.pb(i);
  //    sort(all(indices), [&](const auto& l, const auto& r) {return a[l] <
  //    a[r];}); sort(all(a));

  //    ll num = 0;
  //    for(int i = 1; i <= N; i++)
  //    {
  //        int b = 1, e = N;
  //        while(b <= e)
  //        {
  //            int mid = (b + e) / 2;

  //            if(a[i] + a[mid] == abs(indices[i]-indices[mid]))
  //            {
  //                num++;
  //                break;
  //            }
  //            else if(a[i] + a[mid] > abs(indices[i]-indices[mid]))
  //            {
  //                e = mid - 1;
  //            }
  //            else
  //            {
  //                b = mid + 1;
  //            }
  //        }
  //    }

  //    cout << num << endl;

  return 0;
}


Error: TLE

Fixed submission:

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

#define pii pair<int, int>
#define vi vector<int>
#define pb push_back
#define eb emplace_back
#define mp make_pair
#define x first
#define y second
#define sz(x) (x).size()
#define all(x) (x).begin(), (x).end()  // Forward traversal
#define rall(x) (x).rbegin, (x).rend() // reverse traversal

#define ll long long
#define ld long double

const ld EPS = 10e-9;
const int INF = 1e9 + 5;
const long long INFLL = (ll)INF * (ll)INF;
const ll MOD = 1e9 + 7;
const ld PI = 3.14159265358979323846;

template <typename T> ostream &operator<<(ostream &os, const vector<T> &v) {
  os << "[";
  for (int i = 0; i < (int)v.size(); ++i) {
    os << v[i];
    if (i != (int)v.size() - 1)
      os << ", ";
  }
  os << "]\\n";
  return os;
}

template <typename T> ostream &operator<<(ostream &os, const set<T> &v) {
  os << "{";
  for (const auto &it : v) {
    os << it;
    if (it != *v.rbegin())
      os << ", ";
  }
  os << "}\\n";
  return os;
}

template <typename T>
ostream &operator<<(ostream &os, const unordered_set<T> &v) {
  os << "{";
  for (const auto &it : v) {
    os << it;
    if (it != *v.rbegin())
      os << ", ";
  }
  os << "}\\n";
  return os;
}

template <typename T, typename S>
ostream &operator<<(ostream &os, const map<T, S> &v) {
  os << "{";
  for (const auto &it : v) {
    os << it.first << " : " << it.second;

    if (it != *v.rbegin())
      os << ", ";
  }

  os << "}\\n";
  return os;
}

template <typename T, typename S>
ostream &operator<<(ostream &os, const unordered_map<T, S> &v) {
  os << "{";
  for (const auto &it : v) {
    os << it.first << " : " << it.second;

    if (it != *v.rbegin())
      os << ", ";
  }

  os << "}\\n";
  return os;
}

template <typename T, typename S>
ostream &operator<<(ostream &os, const pair<T, S> &v) {
  os << "(";
  os << v.first << ", " << v.second << ")";
  os << "\\n";
  return os;
}

#define DEBUG

#ifdef DEBUG
#define debug(...) __f(#__VA_ARGS__, __VA_ARGS__)
template <typename Arg1> void __f(const char *name, Arg1 &&arg1) {
  cerr << name << " : " << arg1;
}
template <typename Arg1, typename... Args>
void __f(const char *names, Arg1 &&arg1, Args &&...args) {
  const char *comma = strchr(names + 1, ',');
  cerr.write(names, comma - names) << " : " << arg1 << " | ";
  __f(comma + 1, args...);
}
#else
#define debug(...)
#endif

// __gcd(x, y)

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(0);

  int N;

  cin >> N;

  vector<ll> a(N + 1);
  for (int i = 1; i <= N; i++)
    cin >> a[i];

  unordered_map<ll, ll> s1, s2;

  ll ma = 0;
  for (int i = 1; i <= N; i++) {
    ma = max(ma, i + a[i]);
    s1[i + a[i]]++;
    s2[i - a[i]]++;
  }

  ll num = 0;
  for (int i = 1; i <= ma; i++) {
    if (s1.count(i) && s2.count(i))
      num += s1[i] * s2[i];
  }

  cout << num << endl;

  //    vector<int> indices;
  //    for(int i = 0; i <= N; i++)indices.pb(i);
  //    sort(all(indices), [&](const auto& l, const auto& r) {return a[l] <
  //    a[r];}); sort(all(a));

  //    ll num = 0;
  //    for(int i = 1; i <= N; i++)
  //    {
  //        int b = 1, e = N;
  //        while(b <= e)
  //        {
  //            int mid = (b + e) / 2;

  //            if(a[i] + a[mid] == abs(indices[i]-indices[mid]))
  //            {
  //                num++;
  //                break;
  //            }
  //            else if(a[i] + a[mid] > abs(indices[i]-indices[mid]))
  //            {
  //                e = mid - 1;
  //            }
  //            else
  //            {
  //                b = mid + 1;
  //            }
  //        }
  //    }

  //    cout << num << endl;

  return 0;
}