Problem p03153 (C++)

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>

using namespace std;

#define FOR(i, x, y) for (int i = (x); i < (y); ++i)
#define REP(i, x, y) for (int i = (x); i <= (y); ++i)
#define MP make_pair
#define PB push_back
#define PH push
#define fst first
#define snd second
typedef long long ll;
typedef unsigned long long ull;
typedef double db;
typedef long double ldb;
typedef pair<int, int> pii;

const int maxn = 1e5 + 5;
const ll INF = 1e18;

int n;
ll d, ans;
ll a[maxn], tmp[maxn];
vector<int> vs[maxn];

class SegmentTree {
private:
  static const int siz = 262144;
  pair<ll, int> dat[siz << 1];

public:
  inline int size() { return siz; }
  inline void build(int x, int l, int r) {
    if (l == r) {
      if (l < n)
        dat[x] = MP(tmp[l], l);
      else
        dat[x] = MP(INF, -1);
      return;
    }
    int md = l + r >> 1;
    build(x << 1, l, md);
    build(x << 1 | 1, md + 1, r);
    dat[x] = min(dat[x << 1], dat[x << 1 | 1]);
    return;
  }
  inline void update(int x, int l, int r, int y, ll k) {
    if (l == r) {
      dat[x].fst = k;
      return;
    }
    int md = l + r >> 1;
    if (y <= md)
      update(x << 1, l, md, y, k);
    else
      update(x << 1 | 1, md + 1, r, y, k);
    dat[x] = min(dat[x << 1], dat[x << 1 | 1]);
  }
  inline pair<ll, int> query(int x, int l, int r, int s, int t) {
    if (l >= s && r <= t) {
      return dat[x];
    }
    int md = l + r >> 1;
    if (s > md)
      return query(x << 1 | 1, md + 1, r, s, t);
    if (t <= md)
      return query(x << 1, l, md, s, t);
    return min(query(x << 1, l, md, s, t), query(x << 1 | 1, md + 1, r, s, t));
  }
} seg1, seg2;

class Dsu {
private:
  int fa[maxn], sz[maxn];

public:
  inline void init() {
    FOR(i, 0, n) {
      fa[i] = i;
      sz[i] = 1;
    }
    return;
  }
  inline int find(int x) { return (x == fa[x] ? x : fa[x] = find(fa[x])); }
  inline int size(int x) { return sz[find(x)]; }
  inline bool combine(int x, int y) {
    x = find(x);
    y = find(y);
    if (x == y)
      return false;
    if (sz[x] < sz[y])
      swap(x, y);
    FOR(i, 0, vs[y].size())
    vs[x].PB(vs[y][i]);
    sz[x] += sz[y];
    fa[y] = x;
    return true;
  }
} dsu;

int main() {
  scanf("%d%lld", &n, &d);
  FOR(i, 0, n) { scanf("%lld", a + i); }
  FOR(i, 0, n)
  tmp[i] = -i * d + a[i];
  seg1.build(1, 0, seg1.size() - 1);
  FOR(i, 0, n)
  tmp[i] = i * d + a[i];
  seg2.build(1, 0, seg2.size() - 1);
  dsu.init();
  FOR(i, 0, n)
  vs[i].PB(i);
  while (dsu.size(0) != n) {
    vector<pair<pii, ll>> ed;
    FOR(u, 0, n) if (dsu.find(u) == u) {
      pii e;
      ll mn = INF;
      FOR(i, 0, vs[u].size()) {
        int v = vs[u][i];
        seg1.update(1, 0, seg1.size() - 1, v, INF);
        seg2.update(1, 0, seg2.size() - 1, v, INF);
      }
      FOR(i, 0, vs[u].size()) {
        int v = vs[u][i];
        pair<ll, int> res1 = seg1.query(1, 0, seg1.size() - 1, 0, v);
        pair<ll, int> res2 = seg2.query(1, 0, seg2.size() - 1, v, n - 1);
        if (res1.fst + v * d + a[v] < mn) {
          e = MP(v, res1.snd);
          mn = res1.fst + v * d + a[v];
        }
        if (res2.fst - v * d + a[v] < mn) {
          e = MP(v, res2.snd);
          mn = res2.fst - v * d + a[v];
        }
      }
      if (mn != INF) {
        ed.PB(MP(e, mn));
      }
      FOR(i, 0, vs[u].size()) {
        int v = vs[u][i];
        seg1.update(1, 0, seg1.size() - 1, v, -v * d + a[v]);
        seg2.update(1, 0, seg2.size() - 1, v, v * d + a[v]);
      }
    }
    FOR(i, 0, ed.size())
    ans += (dsu.combine(ed[i].fst.fst, ed[i].fst.snd)) * ed[i].snd;
  }
  printf("%lld\n", ans);
  return 0;
}


Fixed submission:
#include <bits/stdc++.h>

using namespace std;

#define FOR(i, x, y) for (int i = (x); i < (y); ++i)
#define REP(i, x, y) for (int i = (x); i <= (y); ++i)
#define MP make_pair
#define PB push_back
#define PH push
#define fst first
#define snd second
typedef long long ll;
typedef unsigned long long ull;
typedef double db;
typedef long double ldb;
typedef pair<int, int> pii;

const int maxn = 2e5 + 5;
const ll INF = 1e18;

int n;
ll d, ans;
ll a[maxn], tmp[maxn];
vector<int> vs[maxn];

class SegmentTree {
private:
  static const int siz = 262144;
  pair<ll, int> dat[siz << 1];

public:
  inline int size() { return siz; }
  inline void build(int x, int l, int r) {
    if (l == r) {
      if (l < n)
        dat[x] = MP(tmp[l], l);
      else
        dat[x] = MP(INF, -1);
      return;
    }
    int md = l + r >> 1;
    build(x << 1, l, md);
    build(x << 1 | 1, md + 1, r);
    dat[x] = min(dat[x << 1], dat[x << 1 | 1]);
    return;
  }
  inline void update(int x, int l, int r, int y, ll k) {
    if (l == r) {
      dat[x].fst = k;
      return;
    }
    int md = l + r >> 1;
    if (y <= md)
      update(x << 1, l, md, y, k);
    else
      update(x << 1 | 1, md + 1, r, y, k);
    dat[x] = min(dat[x << 1], dat[x << 1 | 1]);
  }
  inline pair<ll, int> query(int x, int l, int r, int s, int t) {
    if (l >= s && r <= t) {
      return dat[x];
    }
    int md = l + r >> 1;
    if (s > md)
      return query(x << 1 | 1, md + 1, r, s, t);
    if (t <= md)
      return query(x << 1, l, md, s, t);
    return min(query(x << 1, l, md, s, t), query(x << 1 | 1, md + 1, r, s, t));
  }
} seg1, seg2;

class Dsu {
private:
  int fa[maxn], sz[maxn];

public:
  inline void init() {
    FOR(i, 0, n) {
      fa[i] = i;
      sz[i] = 1;
    }
    return;
  }
  inline int find(int x) { return (x == fa[x] ? x : fa[x] = find(fa[x])); }
  inline int size(int x) { return sz[find(x)]; }
  inline bool combine(int x, int y) {
    x = find(x);
    y = find(y);
    if (x == y)
      return false;
    if (sz[x] < sz[y])
      swap(x, y);
    FOR(i, 0, vs[y].size())
    vs[x].PB(vs[y][i]);
    sz[x] += sz[y];
    fa[y] = x;
    return true;
  }
} dsu;

int main() {
  scanf("%d%lld", &n, &d);
  FOR(i, 0, n) { scanf("%lld", a + i); }
  FOR(i, 0, n)
  tmp[i] = -i * d + a[i];
  seg1.build(1, 0, seg1.size() - 1);
  FOR(i, 0, n)
  tmp[i] = i * d + a[i];
  seg2.build(1, 0, seg2.size() - 1);
  dsu.init();
  FOR(i, 0, n)
  vs[i].PB(i);
  while (dsu.size(0) != n) {
    vector<pair<pii, ll>> ed;
    FOR(u, 0, n) if (dsu.find(u) == u) {
      pii e;
      ll mn = INF;
      FOR(i, 0, vs[u].size()) {
        int v = vs[u][i];
        seg1.update(1, 0, seg1.size() - 1, v, INF);
        seg2.update(1, 0, seg2.size() - 1, v, INF);
      }
      FOR(i, 0, vs[u].size()) {
        int v = vs[u][i];
        pair<ll, int> res1 = seg1.query(1, 0, seg1.size() - 1, 0, v);
        pair<ll, int> res2 = seg2.query(1, 0, seg2.size() - 1, v, n - 1);
        if (res1.fst + v * d + a[v] < mn) {
          e = MP(v, res1.snd);
          mn = res1.fst + v * d + a[v];
        }
        if (res2.fst - v * d + a[v] < mn) {
          e = MP(v, res2.snd);
          mn = res2.fst - v * d + a[v];
        }
      }
      if (mn != INF) {
        ed.PB(MP(e, mn));
      }
      FOR(i, 0, vs[u].size()) {
        int v = vs[u][i];
        seg1.update(1, 0, seg1.size() - 1, v, -v * d + a[v]);
        seg2.update(1, 0, seg2.size() - 1, v, v * d + a[v]);
      }
    }
    FOR(i, 0, ed.size())
    ans += (dsu.combine(ed[i].fst.fst, ed[i].fst.snd)) * ed[i].snd;
  }
  printf("%lld\n", ans);
  return 0;
}