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

/* Manish Pant - insider_pants, hello_my_friend */

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>

using namespace __gnu_pbds;
using namespace std;

#define F first
#define S second
#define endl "\\n"
#define pb push_back

typedef long long ll;
typedef long double ldob;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;

template <class T> ostream &operator<<(ostream &os, const vector<T> &v) {
  for (T i : v)
    os << i << " ";
  return os;
}
template <class T> ostream &operator<<(ostream &os, const set<T> &v) {
  for (T i : v)
    os << i << " ";
  return os;
}
template <class T> ostream &operator<<(ostream &os, const multiset<T> &v) {
  for (T i : v)
    os << i << " ";
  return os;
}
template <class T, class S>
ostream &operator<<(ostream &os, const pair<T, S> &v) {
  os << v.F << ' ' << v.S;
  return os;
}
template <class T, class S>
ostream &operator<<(ostream &os, const map<T, S> &v) {
  for (auto i : v)
    os << '(' << i.F << " => " << i.S << ')' << ' ';
  return os;
}

void fast() {
  ios_base::sync_with_stdio(false);
  cin.tie(NULL);
  cout.tie(NULL);
}

const int N = 2e5 + 5;
int n;
vector<int> v;
vector<vector<int>> segtree(4 * N);

void build(int node = 0, int l = 0, int r = n - 1) {
  if (l == r) {
    segtree[node].pb(v[l] + l);
    return;
  }
  int mid = (l + r) >> 1;
  build(2 * node + 1, l, mid);
  build(2 * node + 2, mid + 1, r);

  merge(segtree[2 * node + 1].begin(), segtree[2 * node + 1].end(),
        segtree[2 * node + 2].begin(), segtree[2 * node + 2].end(),
        back_inserter(segtree[node]));
}

int query(int x, int y, int val, int node = 0, int l = 0, int r = n - 1) {
  if (l > y || r < x || x > y)
    return 0;
  if (l >= x && r <= y) {
    auto it = upper_bound(segtree[node].begin(), segtree[node].end(), val);
    auto it2 = lower_bound(segtree[node].begin(), segtree[node].end(), val);

    return int(it - it2);
  }
  int mid = (l + r) >> 1;
  return query(x, y, val, 2 * node + 1, l, mid) +
         query(x, y, val, 2 * node + 2, mid + 1, r);
}

void solve() {
  build();
  ll ans = 0;
  for (int i = 0; i < n; i++) {
    int mx = max(i, n - i - 1);
    if (v[i] >= mx)
      continue;

    int diff = mx - v[i];
    ans += query(0, i - v[i] - 1, i - v[i]);

    // cout << i <<"-> ";
    // cout << mx << ' ' << v[i] << endl;
  }
  // ans >>= 1;
  cout << ans << endl;
}

int main() {
  fast();
  cin >> n;
  for (int i = 0; i < n; i++) {
    int a;
    cin >> a;
    v.pb(a);
  }
  solve();
}

Error: TLE

Fixed submission:

/* Manish Pant - insider_pants, hello_my_friend */

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>

using namespace __gnu_pbds;
using namespace std;

#define F first
#define S second
#define endl "\\n"
#define pb push_back

typedef long long ll;
typedef long double ldob;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;

template <class T> ostream &operator<<(ostream &os, const vector<T> &v) {
  for (T i : v)
    os << i << " ";
  return os;
}
template <class T> ostream &operator<<(ostream &os, const set<T> &v) {
  for (T i : v)
    os << i << " ";
  return os;
}
template <class T> ostream &operator<<(ostream &os, const multiset<T> &v) {
  for (T i : v)
    os << i << " ";
  return os;
}
template <class T, class S>
ostream &operator<<(ostream &os, const pair<T, S> &v) {
  os << v.F << ' ' << v.S;
  return os;
}
template <class T, class S>
ostream &operator<<(ostream &os, const map<T, S> &v) {
  for (auto i : v)
    os << '(' << i.F << " => " << i.S << ')' << ' ';
  return os;
}

void fast() {
  ios_base::sync_with_stdio(false);
  cin.tie(NULL);
  cout.tie(NULL);
}

const int N = 2e5 + 5;
int n;
vector<int> v;
vector<vector<int>> segtree(4 * N);

void build(int node = 0, int l = 0, int r = n - 1) {
  if (l == r) {
    segtree[node].pb(v[l] + l);
    return;
  }
  int mid = (l + r) >> 1;
  build(2 * node + 1, l, mid);
  build(2 * node + 2, mid + 1, r);

  merge(segtree[2 * node + 1].begin(), segtree[2 * node + 1].end(),
        segtree[2 * node + 2].begin(), segtree[2 * node + 2].end(),
        back_inserter(segtree[node]));
}

int query(int x, int y, int val, int node = 0, int l = 0, int r = n - 1) {
  if (l > y || r < x || x > y)
    return 0;
  if (l >= x && r <= y) {
    return count(segtree[node].begin(), segtree[node].end(), val);
  }
  int mid = (l + r) >> 1;
  return query(x, y, val, 2 * node + 1, l, mid) +
         query(x, y, val, 2 * node + 2, mid + 1, r);
}

void solve() {
  build();
  ll ans = 0;
  for (int i = 0; i < n; i++) {
    int mx = max(i, n - i - 1);
    if (v[i] >= mx)
      continue;

    int diff = mx - v[i];
    ans += query(0, i - v[i] - 1, i - v[i]);

    // cout << i <<"-> ";
    // cout << mx << ' ' << v[i] << endl;
  }
  // ans >>= 1;
  cout << ans << endl;
}

int main() {
  fast();
  cin >> n;
  for (int i = 0; i < n; i++) {
    int a;
    cin >> a;
    v.pb(a);
  }
  solve();
}