Problem p02691 (C++)

Buggy submission (status: Runtime Error):
/*
    Author: Alam Khan
    AUST CSE 40th Batch

*/

#include <bits/stdc++.h>
using namespace std;
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/detail/standard_policies.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace __gnu_pbds;

typedef long long ll;
typedef pair<ll, ll> pll;
typedef vector<ll> vl;
typedef deque<ll> dl;
typedef stack<ll> stl;
typedef set<ll> sl;
typedef map<string, ll> msl;
typedef map<ll, ll> mll;
typedef tree<ll, null_type, less_equal<ll>, rb_tree_tag,
             tree_order_statistics_node_update>
    ordered_set;
#define sf(n) scanf("%lld", &n)
#define sff(n, m) scanf("%lld %lld", &n, &m)
#define sfff(n, m, r) scanf("%lld %lld %lld", &n, &m, &r)
#define sfs(n) scanf("%s", n)
#define pf(n) printf("%lld\n", n)
#define pff(n, m) printf("%lld %lld\n", n, m)
#define pfff(n, m, r) printf("%lld %lld %lld\n", n, m, r)
#define pfs(n) printf("%s\n", n)
#define pfcs(i, n) printf("Case %lld: %lld\n", i, n)
#define pb push_back
#define prf printf
#define inf 2e18
#define low -1000000000000
#define PI acos(-1.0)
#define rep1(i, n) for (i = 1; i < n; i++)
#define rep0(i, n) for (i = 0; i < n; i++)
#define rep(i, a, n) for (i = a; i < n; i++)
#define repe1(i, n) for (i = 1; i <= n; i++)
#define repe0(i, n) for (i = 0; i <= n; i++)
#define repe(i, a, n) for (i = a; i <= n; i++)
#define endl "\n"
#define FAST                                                                   \
  ios_base::sync_with_stdio(false);                                            \
  cin.tie(NULL);
#define fr freopen("input.txt", "r", stdin)
#define fw freopen("output.txt", "w", stdout)
#define all(v) v.begin(), v.end()
ll a[200009], diff[200009];
vector<ll> v[400009];
void build(ll att, ll L, ll R) {
  if (L == R) {
    v[att].pb(diff[L]);
    return;
  }

  ll mid = (L + R) / 2;
  build(att * 2, L, mid);
  build(att * 2 + 1, mid + 1, R);
  merge(all(v[att * 2]), all(v[att * 2 + 1]), back_inserter(v[att]));
}
ll query(ll att, ll L, ll R, ll l, ll r, ll p) {
  if (l > R || r < L)
    return 0;
  if (l <= L && r >= R) {
    // for(ll i=0;i<v[att].size();i++)
    //   cout<<v[att][i]<<" ";
    // cout<<endl;
    ll x = lower_bound(v[att].begin(), v[att].end(), p) - v[att].begin();
    ll y = upper_bound(v[att].begin(), v[att].end(), p) - v[att].begin();
    return y - x;
  }
  ll mid = (L + R) / 2;
  ll x = query(att * 2, L, mid, l, r, p);
  ll y = query(att * 2 + 1, mid + 1, R, l, r, p);

  return x + y;
}
int main() {
  ll i, n, t, k, j, x = 0, y = 0, m;
  cin >> n;
  for (i = 1; i <= n; i++) {
    cin >> a[i];
  }
  for (i = 1; i <= n; i++) {
    diff[i] = i - a[i];
  }
  build(1, 1, n);
  ll ans = 0;
  for (i = 1; i <= n; i++) {
    x = i + a[i];
    if (x <= n)
      ans += query(1, 1, n, x, n, x);
  }
  cout << ans << endl;
  return 0;
}


Fixed submission:
/*
    Author: Alam Khan
    AUST CSE 40th Batch

*/

#include <bits/stdc++.h>
using namespace std;
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/detail/standard_policies.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace __gnu_pbds;

typedef long long ll;
typedef pair<ll, ll> pll;
typedef vector<ll> vl;
typedef deque<ll> dl;
typedef stack<ll> stl;
typedef set<ll> sl;
typedef map<string, ll> msl;
typedef map<ll, ll> mll;
typedef tree<ll, null_type, less_equal<ll>, rb_tree_tag,
             tree_order_statistics_node_update>
    ordered_set;
#define sf(n) scanf("%lld", &n)
#define sff(n, m) scanf("%lld %lld", &n, &m)
#define sfff(n, m, r) scanf("%lld %lld %lld", &n, &m, &r)
#define sfs(n) scanf("%s", n)
#define pf(n) printf("%lld\n", n)
#define pff(n, m) printf("%lld %lld\n", n, m)
#define pfff(n, m, r) printf("%lld %lld %lld\n", n, m, r)
#define pfs(n) printf("%s\n", n)
#define pfcs(i, n) printf("Case %lld: %lld\n", i, n)
#define pb push_back
#define prf printf
#define inf 2e18
#define low -1000000000000
#define PI acos(-1.0)
#define rep1(i, n) for (i = 1; i < n; i++)
#define rep0(i, n) for (i = 0; i < n; i++)
#define rep(i, a, n) for (i = a; i < n; i++)
#define repe1(i, n) for (i = 1; i <= n; i++)
#define repe0(i, n) for (i = 0; i <= n; i++)
#define repe(i, a, n) for (i = a; i <= n; i++)
#define endl "\n"
#define FAST                                                                   \
  ios_base::sync_with_stdio(false);                                            \
  cin.tie(NULL);
#define fr freopen("input.txt", "r", stdin)
#define fw freopen("output.txt", "w", stdout)
#define all(v) v.begin(), v.end()
ll a[200009], diff[200009];
vector<ll> v[800009];
void build(ll att, ll L, ll R) {
  if (L == R) {
    v[att].pb(diff[L]);
    return;
  }

  ll mid = (L + R) / 2;
  build(att * 2, L, mid);
  build(att * 2 + 1, mid + 1, R);
  merge(all(v[att * 2]), all(v[att * 2 + 1]), back_inserter(v[att]));
}
ll query(ll att, ll L, ll R, ll l, ll r, ll p) {
  if (l > R || r < L)
    return 0;
  if (l <= L && r >= R) {
    // for(ll i=0;i<v[att].size();i++)
    //   cout<<v[att][i]<<" ";
    // cout<<endl;
    ll x = lower_bound(v[att].begin(), v[att].end(), p) - v[att].begin();
    ll y = upper_bound(v[att].begin(), v[att].end(), p) - v[att].begin();
    return y - x;
  }
  ll mid = (L + R) / 2;
  ll x = query(att * 2, L, mid, l, r, p);
  ll y = query(att * 2 + 1, mid + 1, R, l, r, p);

  return x + y;
}
int main() {
  ll i, n, t, k, j, x = 0, y = 0, m;
  cin >> n;
  for (i = 1; i <= n; i++) {
    cin >> a[i];
  }
  for (i = 1; i <= n; i++) {
    diff[i] = i - a[i];
  }
  build(1, 1, n);
  ll ans = 0;
  for (i = 1; i <= n; i++) {
    x = i + a[i];
    if (x <= n)
      ans += query(1, 1, n, x, n, x);
  }
  cout << ans << endl;
  return 0;
}