Problem p03153 (C++)

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>
#define ADD(a, b) a = (a + ll(b)) % mod
#define MUL(a, b) a = (a * ll(b)) % mod
#define MAX(a, b) a = max(a, b)
#define MIN(a, b) a = min(a, b)
#define rep(i, a, b) for (int i = int(a); i < int(b); i++)
#define rer(i, a, b) for (int i = int(a) - 1; i >= int(b); i--)
#define all(a) (a).begin(), (a).end()
#define sz(v) (int)(v).size()
#define pb push_back
#define sec second
#define fst first
#define debug(fmt, ...) Debug(__LINE__, ":", fmt, ##__VA_ARGS__)
using namespace std;
typedef long long ll;
typedef unsigned long long ull;
typedef pair<int, int> pi;
typedef pair<ll, ll> pl;
typedef pair<int, pi> ppi;
typedef vector<int> vi;
typedef vector<ll> vl;
typedef vector<vl> mat;
void Debug() { cout << '\n'; }
template <class FIRST, class... REST> void Debug(FIRST arg, REST... rest) {
  cout << arg << " ";
  Debug(rest...);
}
template <class T> ostream &operator<<(ostream &out, const vector<T> &v) {
  out << "[";
  if (!v.empty()) {
    rep(i, 0, sz(v) - 1) out << v[i] << ", ";
    out << v.back();
  }
  out << "]";
  return out;
}
template <class S, class T>
ostream &operator<<(ostream &out, const pair<S, T> &v) {
  out << "(" << v.first << ", " << v.second << ")";
  return out;
}
const int MAX_N = 300010;
const int MAX_V = 100010;
const double eps = 1e-6;
const ll mod = 1000000007;
const int inf = 1 << 30;
const ll linf = 1LL << 60;
const double PI = 3.14159265358979323846;
///////////////////////////////////////////////////////////////////////////////////////////////////

struct UF {
  vector<int> par, ran;
  void init(int n) {
    par.resize(n);
    ran.resize(n);
    for (int i = 0; i < n; i++) {
      par[i] = i;
      ran[i] = 0;
    }
  }
  UF(int mx = 0) { init(mx); }

  int find(int x) {
    if (par[x] == x)
      return x;
    else
      return par[x] = find(par[x]);
  }
  void unite(int x, int y) {
    x = find(x);
    y = find(y);
    if (x == y)
      return;
    if (ran[x] < ran[y]) {
      par[x] = y;
    } else {
      par[y] = x;
      if (ran[x] == ran[y])
        ran[x]++;
    }
  }
  bool same(int x, int y) { return find(x) == find(y); }
};

//////////////

struct edge {
  int u, v;
  ll cost;
};

bool comp(const edge &e1, const edge &e2) { return e1.cost < e2.cost; }

int N, E;
edge es[MAX_N];

ll kruskal() {
  sort(es, es + E, comp);
  UF uf(N); // init union_find
  ll res = 0;
  for (int i = 0; i < E; i++) {
    edge e = es[i];
    if (!uf.same(e.u, e.v)) {
      uf.unite(e.u, e.v);
      res += e.cost;
    }
  }
  return res;
}

void add_edge(int s, int t, ll cost) {
  // debug(s, t, cost);
  es[E++] = edge{s, t, cost};
}

struct T {
  pl v;
  T(pl v = pl(linf, -1)) : v(v) {}
  T operator+(const T &t) { return T(min(v, t.v)); }
  friend ostream &operator<<(ostream &out, const T &t) {
    out << t.v;
    return out;
  }
};

struct segtree {
  int n;
  vector<T> seg;
  void init(int mx) {
    n = 1;
    while (n < mx)
      n *= 2;
    seg = vector<T>(n * 2);
  }
  segtree(int mx = 0) { init(mx); }
  void update(int i, T x) {
    i += n - 1;
    seg[i] = x;
    while (i > 0) {
      i = (i - 1) / 2;
      seg[i] = seg[i * 2 + 1] + seg[i * 2 + 2];
    }
  }
  T ga(int a, int b, int k, int l, int r) {
    if (b <= l || r <= a)
      return T();
    if (a <= l && r <= b)
      return seg[k];
    else {
      T lv = ga(a, b, k * 2 + 1, l, (l + r) / 2);
      T rv = ga(a, b, k * 2 + 2, (l + r) / 2, r);
      return lv + rv;
    }
  }
  void show() {
    vector<T> tmp;
    rep(i, 0, n) tmp.push_back(get(i, i + 1));
    debug(tmp);
  }
  // edit from here
  T get(int a, int b) { return ga(a, b, 0, 0, n); } //[a, b)
};

ll D;
ll A[MAX_N];
int idx[MAX_N];

void solve() {
  cin >> N >> D;
  rep(i, 0, N) cin >> A[i];
  rep(i, 0, N) idx[i] = i;
  sort(idx, idx + N, [&](int a, int b) { return A[a] < A[b]; });
  // debug(vi(idx, idx + N));
  segtree seg1(N); // Ai - D * i
  segtree seg2(N); // Ai + D * i

  rep(j, 0, N) {
    int i = idx[j];
    pl lp = seg1.get(0, i).v;
    pl rp = seg2.get(i + 1, N).v;
    if (lp.fst != linf) {
      add_edge(i, lp.sec, A[lp.sec] + A[i] + D * (i - lp.sec));
    }
    if (rp.fst != linf) {
      add_edge(i, rp.sec, A[i] + A[rp.sec] + D * (rp.sec - i));
    }
    seg1.update(i, T(pl(A[i] - D * i, i)));
    seg2.update(i, T(pl(A[i] + D * i, i)));
  }
  cout << kruskal() << "\n";
}

int main() {
#ifndef LOCAL
  ios::sync_with_stdio(false);
  cin.tie(0);
#endif
  cout << fixed;
  cout.precision(20);
  srand((unsigned int)time(NULL));
#ifdef LOCAL
  // freopen("in.txt", "wt", stdout); //for tester
  freopen("in.txt", "rt", stdin);
#endif
  solve();
#ifdef LOCAL
  cerr << "Time elapsed: " << 1.0 * clock() / CLOCKS_PER_SEC << " s.\n";
#endif
  return 0;
}


Fixed submission:
#include <bits/stdc++.h>
#define ADD(a, b) a = (a + ll(b)) % mod
#define MUL(a, b) a = (a * ll(b)) % mod
#define MAX(a, b) a = max(a, b)
#define MIN(a, b) a = min(a, b)
#define rep(i, a, b) for (int i = int(a); i < int(b); i++)
#define rer(i, a, b) for (int i = int(a) - 1; i >= int(b); i--)
#define all(a) (a).begin(), (a).end()
#define sz(v) (int)(v).size()
#define pb push_back
#define sec second
#define fst first
#define debug(fmt, ...) Debug(__LINE__, ":", fmt, ##__VA_ARGS__)
using namespace std;
typedef long long ll;
typedef unsigned long long ull;
typedef pair<int, int> pi;
typedef pair<ll, ll> pl;
typedef pair<int, pi> ppi;
typedef vector<int> vi;
typedef vector<ll> vl;
typedef vector<vl> mat;
void Debug() { cout << '\n'; }
template <class FIRST, class... REST> void Debug(FIRST arg, REST... rest) {
  cout << arg << " ";
  Debug(rest...);
}
template <class T> ostream &operator<<(ostream &out, const vector<T> &v) {
  out << "[";
  if (!v.empty()) {
    rep(i, 0, sz(v) - 1) out << v[i] << ", ";
    out << v.back();
  }
  out << "]";
  return out;
}
template <class S, class T>
ostream &operator<<(ostream &out, const pair<S, T> &v) {
  out << "(" << v.first << ", " << v.second << ")";
  return out;
}
const int MAX_N = 300010;
const int MAX_V = 100010;
const double eps = 1e-6;
const ll mod = 1000000007;
const int inf = 1 << 30;
const ll linf = 1LL << 60;
const double PI = 3.14159265358979323846;
///////////////////////////////////////////////////////////////////////////////////////////////////

struct UF {
  vector<int> par, ran;
  void init(int n) {
    par.resize(n);
    ran.resize(n);
    for (int i = 0; i < n; i++) {
      par[i] = i;
      ran[i] = 0;
    }
  }
  UF(int mx = 0) { init(mx); }

  int find(int x) {
    if (par[x] == x)
      return x;
    else
      return par[x] = find(par[x]);
  }
  void unite(int x, int y) {
    x = find(x);
    y = find(y);
    if (x == y)
      return;
    if (ran[x] < ran[y]) {
      par[x] = y;
    } else {
      par[y] = x;
      if (ran[x] == ran[y])
        ran[x]++;
    }
  }
  bool same(int x, int y) { return find(x) == find(y); }
};

//////////////

struct edge {
  int u, v;
  ll cost;
};

bool comp(const edge &e1, const edge &e2) { return e1.cost < e2.cost; }

int N, E;
edge es[MAX_N * 2];

ll kruskal() {
  sort(es, es + E, comp);
  UF uf(N); // init union_find
  ll res = 0;
  for (int i = 0; i < E; i++) {
    edge e = es[i];
    if (!uf.same(e.u, e.v)) {
      uf.unite(e.u, e.v);
      res += e.cost;
    }
  }
  return res;
}

void add_edge(int s, int t, ll cost) {
  // debug(s, t, cost);
  es[E++] = edge{s, t, cost};
}

struct T {
  pl v;
  T(pl v = pl(linf, -1)) : v(v) {}
  T operator+(const T &t) { return T(min(v, t.v)); }
  friend ostream &operator<<(ostream &out, const T &t) {
    out << t.v;
    return out;
  }
};

struct segtree {
  int n;
  vector<T> seg;
  void init(int mx) {
    n = 1;
    while (n < mx)
      n *= 2;
    seg = vector<T>(n * 2);
  }
  segtree(int mx = 0) { init(mx); }
  void update(int i, T x) {
    i += n - 1;
    seg[i] = x;
    while (i > 0) {
      i = (i - 1) / 2;
      seg[i] = seg[i * 2 + 1] + seg[i * 2 + 2];
    }
  }
  T ga(int a, int b, int k, int l, int r) {
    if (b <= l || r <= a)
      return T();
    if (a <= l && r <= b)
      return seg[k];
    else {
      T lv = ga(a, b, k * 2 + 1, l, (l + r) / 2);
      T rv = ga(a, b, k * 2 + 2, (l + r) / 2, r);
      return lv + rv;
    }
  }
  void show() {
    vector<T> tmp;
    rep(i, 0, n) tmp.push_back(get(i, i + 1));
    debug(tmp);
  }
  // edit from here
  T get(int a, int b) { return ga(a, b, 0, 0, n); } //[a, b)
};

ll D;
ll A[MAX_N];
int idx[MAX_N];

void solve() {
  cin >> N >> D;
  rep(i, 0, N) cin >> A[i];
  rep(i, 0, N) idx[i] = i;
  sort(idx, idx + N, [&](int a, int b) { return A[a] < A[b]; });
  // debug(vi(idx, idx + N));
  segtree seg1(N); // Ai - D * i
  segtree seg2(N); // Ai + D * i

  rep(j, 0, N) {
    int i = idx[j];
    pl lp = seg1.get(0, i).v;
    pl rp = seg2.get(i + 1, N).v;
    if (lp.fst != linf) {
      add_edge(i, lp.sec, A[lp.sec] + A[i] + D * (i - lp.sec));
    }
    if (rp.fst != linf) {
      add_edge(i, rp.sec, A[i] + A[rp.sec] + D * (rp.sec - i));
    }
    seg1.update(i, T(pl(A[i] - D * i, i)));
    seg2.update(i, T(pl(A[i] + D * i, i)));
  }
  cout << kruskal() << "\n";
}

int main() {
#ifndef LOCAL
  ios::sync_with_stdio(false);
  cin.tie(0);
#endif
  cout << fixed;
  cout.precision(20);
  srand((unsigned int)time(NULL));
#ifdef LOCAL
  // freopen("in.txt", "wt", stdout); //for tester
  freopen("in.txt", "rt", stdin);
#endif
  solve();
#ifdef LOCAL
  cerr << "Time elapsed: " << 1.0 * clock() / CLOCKS_PER_SEC << " s.\n";
#endif
  return 0;
}