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

// https://atcoder.jp/contests/abc166/tasks/abc166_e
#include <algorithm>
#include <cassert>
#include <chrono>
#include <cstdint>
#include <deque>
#include <functional>
#include <iomanip>
#include <iostream>
#include <limits>
#include <map>
#include <math.h>
#include <numeric>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <vector>
#define REP(i, n) for (ll i = 0, i##_len = (n); i < i##_len; ++i)
#define RREP(i, n) for (ll i = (n)-1; i >= 0; --i)
#define SREP(i, s, n) for (ll i = (s), i##_len = (n); i < i##_len; ++i)
#define ALL(v) (v).begin(), (v).end()

using namespace std;
typedef long long int ll;
typedef unsigned long long int ull;

template <typename T> T gcd(T a, T b) {
  if (a < b)
    gcd(b, a);
  if (b == 1)
    return 1;
  T r;
  while ((r = a % b)) {
    a = b;
    b = r;
  }
  return b;
}

template <typename T> struct UnionFind {
  vector<T> par;
  vector<T> rank;

  UnionFind(T n = 1) { init(n); }

  void init(T n = 1) {
    par.resize(n);
    rank.resize(n);
    for (T i = 0; i < n; ++i) {
      par[i] = i;
      rank[i] = 0;
    }
  }

  int root(T x) {
    if (par[x] == x) {
      return x;
    } else {
      int r = root(par[x]);
      return par[x] = r;
    }
  }

  bool issame(T x, T y) { return root(x) == root(y); }

  bool merge(T x, T y) {
    x = root(x);
    y = root(y);
    if (x == y)
      return false;
    if (rank[x] < rank[y])
      swap(x, y);
    if (rank[x] == rank[y])
      ++rank[x];
    par[y] = x;
    return true;
  }
};

// #define INF (ll(1e18))

template <typename Monoid, typename OperatorMonoid = Monoid>
struct LazySegmentTree {
  using F = function<Monoid(Monoid, Monoid)>;
  using G = function<Monoid(Monoid, OperatorMonoid)>;
  using H = function<OperatorMonoid(OperatorMonoid, OperatorMonoid)>;

  ll sz, height;
  vector<Monoid> data;
  vector<OperatorMonoid> lazy;
  const F f;
  const G g;
  const H h;
  const Monoid M1;
  const OperatorMonoid OM0;

  LazySegmentTree(ll n, const F f, const G g, const H h, const Monoid &M1,
                  const OperatorMonoid OM0)
      : f(f), g(g), h(h), M1(M1), OM0(OM0) {
    sz = 1;
    height = 0;
    while (sz < n)
      sz <<= 1, height++;
    data.assign(2 * sz, M1);
    lazy.assign(2 * sz, OM0);
  }

  void set(ll k, const Monoid &x) { data[k + sz] = x; }

  void build() {
    for (int k = sz - 1; k > 0; k--) {
      data[k] = f(data[2 * k + 0], data[2 * k + 1]);
    }
  }

  inline void propagate(ll k) {
    if (lazy[k] != OM0) {
      lazy[2 * k + 0] = h(lazy[2 * k + 0], lazy[k]);
      lazy[2 * k + 1] = h(lazy[2 * k + 1], lazy[k]);
      data[k] = reflect(k);
      lazy[k] = OM0;
    }
  }

  inline Monoid reflect(ll k) {
    return lazy[k] == OM0 ? data[k] : g(data[k], lazy[k]);
  }

  inline void recalc(ll k) {
    while (k >>= 1)
      data[k] = f(reflect(2 * k + 0), reflect(2 * k + 1));
  }

  inline void thrust(ll k) {
    for (int i = height; i > 0; i--)
      propagate(k >> i);
  }

  void update(ll a, ll b, const OperatorMonoid &x) {
    thrust(a += sz);
    thrust(b += sz - 1);
    for (int l = a, r = b + 1; l < r; l >>= 1, r >>= 1) {
      if (l & 1)
        lazy[l] = h(lazy[l], x), ++l;
      if (r & 1)
        --r, lazy[r] = h(lazy[r], x);
    }
    recalc(a);
    recalc(b);
  }

  Monoid query(ll a, ll b) {
    thrust(a += sz);
    thrust(b += sz - 1);
    Monoid L = M1, R = M1;
    for (int l = a, r = b + 1; l < r; l >>= 1, r >>= 1) {
      if (l & 1)
        L = f(L, reflect(l++));
      if (r & 1)
        R = f(reflect(--r), R);
    }
    return f(L, R);
  }

  Monoid operator[](const int &k) { return query(k, k + 1); }

  template <typename C>
  int find_subtree(ll a, const C &check, Monoid &M, bool type) {
    while (a < sz) {
      propagate(a);
      Monoid nxt =
          type ? f(reflect(2 * a + type), M) : f(M, reflect(2 * a + type));
      if (check(nxt))
        a = 2 * a + type;
      else
        M = nxt, a = 2 * a + 1 - type;
    }
    return a - sz;
  }

  template <typename C> int find_first(ll a, const C &check) {
    Monoid L = M1;
    if (a <= 0) {
      if (check(f(L, reflect(1))))
        return find_subtree(1, check, L, false);
      return -1;
    }
    thrust(a + sz);
    int b = sz;
    for (a += sz, b += sz; a < b; a >>= 1, b >>= 1) {
      if (a & 1) {
        Monoid nxt = f(L, reflect(a));
        if (check(nxt))
          return find_subtree(a, check, L, false);
        L = nxt;
        ++a;
      }
    }
    return -1;
  }

  template <typename C> int find_last(ll b, const C &check) {
    Monoid R = M1;
    if (b >= sz) {
      if (check(f(reflect(1), R)))
        return find_subtree(1, check, R, true);
      return -1;
    }
    thrust(b + sz - 1);
    int a = sz;
    for (b += sz; a < b; a >>= 1, b >>= 1) {
      if (b & 1) {
        Monoid nxt = f(reflect(--b), R);
        if (check(nxt))
          return find_subtree(b, check, R, true);
        R = nxt;
      }
    }
    return -1;
  }
};

int msb(ll x) {
  int n = 0;
  while (x > 0) {
    x /= 2;
    n++;
  }
  return n;
}

template <ll Modulus> class modint {
public:
  ll a;

  constexpr modint(const ll x = 0) noexcept : a(x % Modulus) {}
  constexpr ll &value() noexcept { return a; }
  constexpr operator int() const { return a; }
  constexpr operator long long() const { return a; }
  constexpr const ll &value() const noexcept { return a; }
  constexpr modint operator+(const modint rhs) const noexcept {
    return modint(*this) += rhs;
  }
  constexpr modint operator+(const ll rhs) const noexcept {
    return modint(*this) += rhs;
  }
  constexpr modint operator-(const modint rhs) const noexcept {
    return modint(*this) -= rhs;
  }
  constexpr modint operator-(const ll rhs) const noexcept {
    return modint(*this) -= rhs;
  }
  constexpr modint operator*(const modint rhs) const noexcept {
    return modint(*this) *= rhs;
  }
  constexpr modint operator*(const ll rhs) const noexcept {
    return modint(*this) *= rhs;
  }
  constexpr modint operator/(const modint rhs) const noexcept {
    return modint(*this) /= rhs;
  }
  constexpr modint operator/(const ll rhs) const noexcept {
    return modint(*this) /= rhs;
  }
  constexpr modint &operator+=(const modint rhs) noexcept {
    a += rhs.a;
    if (a >= Modulus) {
      a -= Modulus;
    }
    return *this;
  }
  constexpr modint &operator-=(const modint rhs) noexcept {
    if (a < rhs.a) {
      a += Modulus;
    }
    a -= rhs.a;
    return *this;
  }
  constexpr modint &operator*=(const modint rhs) noexcept {
    a = a * rhs.a % Modulus;
    return *this;
  }
  constexpr modint &operator/=(modint rhs) noexcept {
    ll exp = Modulus - 2;
    while (exp) {
      if (exp % 2) {
        *this *= rhs;
      }
      rhs *= rhs;
      exp /= 2;
    }
    return *this;
  }
};

#define MOD (ll(998244353))
#define INF (ll(1e15))

int main(int argc, const char *argv[]) {
  (void)argc;
  (void)argv;
  cin.tie(0);
  ios::sync_with_stdio(false);
  ll N;
  cin >> N;
  vector<ll> A(N);
  REP(i, N) { cin >> A[i]; }
  map<ll, ll> s;
  ll score = 0;
  REP(i, N) {
    score += s[A[i] - i];
    s[-i - A[i]]++;
  }
  cout << score << endl;
}


Error: TLE

Fixed submission:

// https://atcoder.jp/contests/abc166/tasks/abc166_e
#include <algorithm>
#include <cassert>
#include <chrono>
#include <cstdint>
#include <deque>
#include <functional>
#include <iomanip>
#include <iostream>
#include <limits>
#include <map>
#include <math.h>
#include <numeric>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <vector>
#define REP(i, n) for (ll i = 0, i##_len = (n); i < i##_len; ++i)
#define RREP(i, n) for (ll i = (n)-1; i >= 0; --i)
#define SREP(i, s, n) for (ll i = (s), i##_len = (n); i < i##_len; ++i)
#define ALL(v) (v).begin(), (v).end()

using namespace std;
typedef long long int ll;
typedef unsigned long long int ull;

template <typename T> T gcd(T a, T b) {
  if (a < b)
    gcd(b, a);
  if (b == 1)
    return 1;
  T r;
  while ((r = a % b)) {
    a = b;
    b = r;
  }
  return b;
}

template <typename T> struct UnionFind {
  vector<T> par;
  vector<T> rank;

  UnionFind(T n = 1) { init(n); }

  void init(T n = 1) {
    par.resize(n);
    rank.resize(n);
    for (T i = 0; i < n; ++i) {
      par[i] = i;
      rank[i] = 0;
    }
  }

  int root(T x) {
    if (par[x] == x) {
      return x;
    } else {
      int r = root(par[x]);
      return par[x] = r;
    }
  }

  bool issame(T x, T y) { return root(x) == root(y); }

  bool merge(T x, T y) {
    x = root(x);
    y = root(y);
    if (x == y)
      return false;
    if (rank[x] < rank[y])
      swap(x, y);
    if (rank[x] == rank[y])
      ++rank[x];
    par[y] = x;
    return true;
  }
};

// #define INF (ll(1e18))

template <typename Monoid, typename OperatorMonoid = Monoid>
struct LazySegmentTree {
  using F = function<Monoid(Monoid, Monoid)>;
  using G = function<Monoid(Monoid, OperatorMonoid)>;
  using H = function<OperatorMonoid(OperatorMonoid, OperatorMonoid)>;

  ll sz, height;
  vector<Monoid> data;
  vector<OperatorMonoid> lazy;
  const F f;
  const G g;
  const H h;
  const Monoid M1;
  const OperatorMonoid OM0;

  LazySegmentTree(ll n, const F f, const G g, const H h, const Monoid &M1,
                  const OperatorMonoid OM0)
      : f(f), g(g), h(h), M1(M1), OM0(OM0) {
    sz = 1;
    height = 0;
    while (sz < n)
      sz <<= 1, height++;
    data.assign(2 * sz, M1);
    lazy.assign(2 * sz, OM0);
  }

  void set(ll k, const Monoid &x) { data[k + sz] = x; }

  void build() {
    for (int k = sz - 1; k > 0; k--) {
      data[k] = f(data[2 * k + 0], data[2 * k + 1]);
    }
  }

  inline void propagate(ll k) {
    if (lazy[k] != OM0) {
      lazy[2 * k + 0] = h(lazy[2 * k + 0], lazy[k]);
      lazy[2 * k + 1] = h(lazy[2 * k + 1], lazy[k]);
      data[k] = reflect(k);
      lazy[k] = OM0;
    }
  }

  inline Monoid reflect(ll k) {
    return lazy[k] == OM0 ? data[k] : g(data[k], lazy[k]);
  }

  inline void recalc(ll k) {
    while (k >>= 1)
      data[k] = f(reflect(2 * k + 0), reflect(2 * k + 1));
  }

  inline void thrust(ll k) {
    for (int i = height; i > 0; i--)
      propagate(k >> i);
  }

  void update(ll a, ll b, const OperatorMonoid &x) {
    thrust(a += sz);
    thrust(b += sz - 1);
    for (int l = a, r = b + 1; l < r; l >>= 1, r >>= 1) {
      if (l & 1)
        lazy[l] = h(lazy[l], x), ++l;
      if (r & 1)
        --r, lazy[r] = h(lazy[r], x);
    }
    recalc(a);
    recalc(b);
  }

  Monoid query(ll a, ll b) {
    thrust(a += sz);
    thrust(b += sz - 1);
    Monoid L = M1, R = M1;
    for (int l = a, r = b + 1; l < r; l >>= 1, r >>= 1) {
      if (l & 1)
        L = f(L, reflect(l++));
      if (r & 1)
        R = f(reflect(--r), R);
    }
    return f(L, R);
  }

  Monoid operator[](const int &k) { return query(k, k + 1); }

  template <typename C>
  int find_subtree(ll a, const C &check, Monoid &M, bool type) {
    while (a < sz) {
      propagate(a);
      Monoid nxt =
          type ? f(reflect(2 * a + type), M) : f(M, reflect(2 * a + type));
      if (check(nxt))
        a = 2 * a + type;
      else
        M = nxt, a = 2 * a + 1 - type;
    }
    return a - sz;
  }

  template <typename C> int find_first(ll a, const C &check) {
    Monoid L = M1;
    if (a <= 0) {
      if (check(f(L, reflect(1))))
        return find_subtree(1, check, L, false);
      return -1;
    }
    thrust(a + sz);
    int b = sz;
    for (a += sz, b += sz; a < b; a >>= 1, b >>= 1) {
      if (a & 1) {
        Monoid nxt = f(L, reflect(a));
        if (check(nxt))
          return find_subtree(a, check, L, false);
        L = nxt;
        ++a;
      }
    }
    return -1;
  }

  template <typename C> int find_last(ll b, const C &check) {
    Monoid R = M1;
    if (b >= sz) {
      if (check(f(reflect(1), R)))
        return find_subtree(1, check, R, true);
      return -1;
    }
    thrust(b + sz - 1);
    int a = sz;
    for (b += sz; a < b; a >>= 1, b >>= 1) {
      if (b & 1) {
        Monoid nxt = f(reflect(--b), R);
        if (check(nxt))
          return find_subtree(b, check, R, true);
        R = nxt;
      }
    }
    return -1;
  }
};

int msb(ll x) {
  int n = 0;
  while (x > 0) {
    x /= 2;
    n++;
  }
  return n;
}

template <ll Modulus> class modint {
public:
  ll a;

  constexpr modint(const ll x = 0) noexcept : a(x % Modulus) {}
  constexpr ll &value() noexcept { return a; }
  constexpr operator int() const { return a; }
  constexpr operator long long() const { return a; }
  constexpr const ll &value() const noexcept { return a; }
  constexpr modint operator+(const modint rhs) const noexcept {
    return modint(*this) += rhs;
  }
  constexpr modint operator+(const ll rhs) const noexcept {
    return modint(*this) += rhs;
  }
  constexpr modint operator-(const modint rhs) const noexcept {
    return modint(*this) -= rhs;
  }
  constexpr modint operator-(const ll rhs) const noexcept {
    return modint(*this) -= rhs;
  }
  constexpr modint operator*(const modint rhs) const noexcept {
    return modint(*this) *= rhs;
  }
  constexpr modint operator*(const ll rhs) const noexcept {
    return modint(*this) *= rhs;
  }
  constexpr modint operator/(const modint rhs) const noexcept {
    return modint(*this) /= rhs;
  }
  constexpr modint operator/(const ll rhs) const noexcept {
    return modint(*this) /= rhs;
  }
  constexpr modint &operator+=(const modint rhs) noexcept {
    a += rhs.a;
    if (a >= Modulus) {
      a -= Modulus;
    }
    return *this;
  }
  constexpr modint &operator-=(const modint rhs) noexcept {
    if (a < rhs.a) {
      a += Modulus;
    }
    a -= rhs.a;
    return *this;
  }
  constexpr modint &operator*=(const modint rhs) noexcept {
    a = a * rhs.a % Modulus;
    return *this;
  }
  constexpr modint &operator/=(modint rhs) noexcept {
    ll exp = Modulus - 2;
    while (exp) {
      if (exp % 2) {
        *this *= rhs;
      }
      rhs *= rhs;
      exp /= 2;
    }
    return *this;
  }
};

#define MOD (ll(998244353))
#define INF (ll(1e15))

int main(int argc, const char *argv[]) {
  (void)argc;
  (void)argv;
  cin.tie(0);
  ios::sync_with_stdio(false);
  ll N;
  cin >> N;
  vector<ll> A(N);
  REP(i, N) { cin >> A[i]; }
  multiset<ll> s;
  ll score = 0;
  REP(i, N) {
    if (s.find(A[i] - i) != s.end()) {
      score += s.count(A[i] - i);
    }
    s.insert(-i - A[i]);
  }
  cout << score << endl;
}