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

#include <algorithm>
#include <cmath>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <map>
#include <numeric>
#include <string>
#include <vector>

typedef long double ld;
typedef long long ll;
const ll INF = (ll)1e18 + 1;
const ll MOD = 1e9 + 7;

// Output
template <class T, class S>
std::ostream &operator<<(std::ostream &os, std::pair<T, S> p) {
  return os << "(" << p.first << " " << p.second << ")";
}

template <class T>
std::ostream &operator<<(std::ostream &os, std::vector<T> v) {
  for (ll i = 0; i < (ll)v.size(); i++) {
    os << " [" << i << "]" << v[i];
    if (i % 10 == 9)
      os << std::endl;
  }
  // os << v[0]; for (ll i = 1; i < (ll)v.size(); i++){ os << " " << v[i]; }
  return os;
}

template <class T, class S>
std::ostream &operator<<(std::ostream &os, std::map<T, S> m) {
  ll i = 0;
  for (auto p : m) {
    os << " [" << i << "]" << p.first << "->" << p.second;
    i++;
  }
  return os;
}

template <class T>
std::ostream &operator<<(std::ostream &os, std::vector<std::vector<T>> vv) {
  for (auto &vec : vv) {
    os << vec << std::endl;
  }
  os << "--------" << std::endl;
  return os;
}

void print() { std::cout << std::endl; }
template <typename H> void print(H &&head) { std::cout << head << std::endl; }
template <typename H, typename... T> void print(H &&head, T &&...tail) {
  std::cout << head << " ", print(std::forward<T>(tail)...);
}

void solve() {
  ll N;
  std::cin >> N;
  std::vector<ll> A(N, 0LL);
  for (ll i = 0; i < N; i++)
    std::cin >> A[i];

  // i - j = A[i] + A[j] なので、
  // i - A[i] = j + A[j] となる
  // まずは、i-A[i] となる値の個数を map を用いて列挙する
  std::map<ll, ll> mcnt;
  for (ll i = 0; i < N; i++) {
    mcnt[i - A[i]]++;
  }

  // i - A[i] == i + A[i] となる場合
  // 0 == 2 * A[i] なので、
  // 制約条件である A[i] >= 1 に該当しなくなる。
  // そのため i - A[i] == i + A[i] にはならない。
  // よって、ある i で i + A[i] を求め、j - A[j]
  // となる個数がいくつあるか調べる。
  ll ans = 0LL;
  for (ll i = 0; i < N; i++) {
    ans += mcnt[i + A[i]];
  }

  print(ans);
}

int main() {
  std::cin.tie(nullptr);
  std::cout.tie(nullptr);
  std::ios::sync_with_stdio(false);

  solve();

  return 0;
}


Error: TLE

Fixed submission:

#include <algorithm>
#include <cmath>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <map>
#include <numeric>
#include <string>
#include <vector>

typedef long double ld;
typedef long long ll;
const ll INF = (ll)1e18 + 1;
const ll MOD = 1e9 + 7;

// Split
namespace util {

std::vector<std::string> split(std::string s, char delimiter) {
  std::vector<std::string> vs;
  std::string sub;
  for (auto c : s) {
    if (c == delimiter)
      vs.push_back(sub), sub.clear();
    else
      sub += c;
  }
  vs.push_back(sub);
  return vs;
}

} // namespace util

// Minimum, Maximum
template <class T> T minimum(T head, T tail) { return std::min(head, tail); }
template <class H, class... T> H minimum(H head, T... tail) {
  return std::min(head, minimum(tail...));
}
template <class T> T maximum(T head, T tail) { return std::max(head, tail); }
template <class H, class... T> H maximum(H head, T... tail) {
  return std::max(head, maximum(tail...));
}

// Output
template <class T, class S>
std::ostream &operator<<(std::ostream &os, std::pair<T, S> p) {
  return os << "(" << p.first << " " << p.second << ")";
}

template <class T>
std::ostream &operator<<(std::ostream &os, std::vector<T> v) {
  for (ll i = 0; i < (ll)v.size(); i++) {
    os << " [" << i << "]" << v[i];
    if (i % 10 == 9)
      os << std::endl;
  }
  // os << v[0]; for (ll i = 1; i < (ll)v.size(); i++){ os << " " << v[i]; }
  return os;
}

template <class T, class S>
std::ostream &operator<<(std::ostream &os, std::map<T, S> m) {
  ll i = 0;
  for (auto p : m) {
    os << " [" << i << "]" << p.first << "->" << p.second;
    i++;
  }
  return os;
}

template <class T>
std::ostream &operator<<(std::ostream &os, std::vector<std::vector<T>> vv) {
  for (auto &vec : vv) {
    os << vec << std::endl;
  }
  os << "--------" << std::endl;
  return os;
}

void print() { std::cout << std::endl; }
template <typename H> void print(H &&head) { std::cout << head << std::endl; }
template <typename H, typename... T> void print(H &&head, T &&...tail) {
  std::cout << head << " ", print(std::forward<T>(tail)...);
}

// GCD, LCM
ll gcd(ll a, ll b) {
  while (b != 0) {
    a %= b;
    std::swap(a, b);
  }
  return a;
}
ll lcm(ll a, ll b) { return (a / gcd(a, b)) * b; }

// DIGNUM
std::vector<ll> dignum(ll x) {
  if (x == 0)
    return {0};
  std::vector<ll> res;
  while (x) {
    res.push_back(x % 10);
    x /= 10;
  }
  std::reverse(res.begin(), res.end());
  return res;
}

// FAIL
void fail() {
  print(-1);
  std::exit(0);
}
// void fail(){ print("No"); std::exit(0); }

int main() {
  std::cin.tie(nullptr);
  std::cout.tie(nullptr);
  std::ios::sync_with_stdio(false);

  ll N;
  std::cin >> N;
  std::vector<ll> va(N, 0LL);
  for (ll i = 0; i < N; i++) {
    std::cin >> va[i];
  }

  ll cnt = 0LL;
  for (ll i = 0; i < N; i++) {
    for (ll j = i + 1; j < N; j++) {
      if (j - i == va[j] + va[i])
        cnt++;
    }
  }

  print(cnt);
  return 0;
}