Problem p03153 (C++)

Buggy submission (status: Runtime Error):
#include <algorithm>
#include <bits/stdc++.h>
#include <bitset>
#include <cassert>
#include <iomanip>
#include <iostream>
#include <map>
#include <math.h>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <vector>
using namespace std;
#define rep(i, n) for (ll i = 0; i < (ll)n; i++)
#define repf(i, a, b) for (ll i = (ll)a; i < (ll)b; i++)
#define repr(i, a, b) for (ll i = (ll)a; i > (ll)b; i--)
#define repv(x, arr) for (auto &x : arr)
#define all(v) (v).begin(), (v).end()
#define vec(name, num) vector<ll> name((num), 0);
#define mp(a, b) make_pair(a, b)
#define pb(x) push_back(x)
#define F first
#define S second
typedef long long ll;
typedef long double ld;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;
typedef pair<ld, ld> pdd;
typedef vector<ll> vll;
typedef vector<vll> vvll;
typedef vector<int> vii;
typedef vector<vii> vvii;
const ll mod = 1e9 + 7;
const int infi = 2147483600;
const ll infl = 1e17;

struct SegmentTree {
private:
  int N;
  vector<ll> node;
  ll e;

public:
  SegmentTree() {}
  SegmentTree(ll e) : e(e) {}
  void __init__(int sz) {
    N = 1;
    while (N < sz)
      N *= 2;
    node.resize(2 * N - 1, e);
  }
  void build(vector<ll> &v) {
    int sz = (int)(v.size());
    __init__(sz);
    rep(i, sz) { node[i + N - 1] = v[i]; }
    for (int i = N - 2; i >= 0; i--) {
      node[i] = min(node[i * 2 + 1], node[i * 2 + 2]);
    }
  }
  void update(int k, ll x) {
    k += N - 1;
    node[k] = x;
    while (k > 0) {
      k = (k - 1) / 2;
      node[k] = min(node[2 * k + 1], node[2 * k + 2]);
    }
    return;
  }
  ll query(int a, int b) { return query(a, b, 0, 0, N); }
  ll query(int a, int b, int k, int l, int r) {
    if (b <= l || r <= a)
      return e;
    if (a <= l && r <= b)
      return node[k];
    ll vl, vr;
    vl = query(a, b, 2 * k + 1, l, (l + r) / 2);
    vr = query(a, b, 2 * k + 2, (l + r) / 2, r);
    return min(vl, vr);
  }
  ll show_one(int i) {
    if (0 <= i && i < N)
      return node[i + N - 1];
    return e;
  }

  int need() {
    int k = 0;
    while (k < N - 1) {
      node[k] = min(node[2 * k + 1], node[2 * k + 2]);
      if (node[k] == node[2 * k + 1]) {
        k = 2 * k + 1;
      }
      if (node[k] == node[2 * k + 2]) {
        k = 2 * k + 2;
      }
    }
    return k - N + 1;
  }
  void view() {
    repf(k, N - 1, 2 * N - 1) {
      if (node[k] == e)
        cout << "inf ";
      else
        cout << node[k] << " ";
    }
    cout << endl;
    return;
  }
};

struct UnionFind {
  vector<int> root;
  int n;
  void __init__(int sz) {
    n = sz;
    root.resize(n, -1);
  }
  int find(int x) {
    if (root[x] >= 0) {
      root[x] = find(root[x]);
      return root[x];
    }
    return x;
  }
  void unit(int a, int b) {
    int ra, rb;
    ra = find(a);
    rb = find(b);
    if (ra == rb) {
      return;
    } else if (root[ra] <= root[rb]) {
      root[ra] += root[rb];
      root[rb] = ra;
    } else {
      root[rb] += root[ra];
      root[ra] = rb;
    }
    return;
  }
  int count() {
    set<int> cnt;
    rep(i, n) cnt.insert(find(i));
    return cnt.size();
  }
};

int main() {
  ll n, d;
  cin >> n >> d;
  vll a(n);
  rep(i, n) cin >> a[i];
  SegmentTree seg(infl);
  seg.build(a);
  ll ans = 0;
  ll cnt = n - 1;
  vvll idx(n, vll(2));
  rep(i, n) {
    idx[i][0] = i;
    idx[i][1] = i;
  }
  while (cnt > 0) {
    /*cout << ans;
    cout << endl;
    seg.view(); */
    int k = seg.need();
    ll x = seg.show_one(k);
    // cout << k << " " << x << " ";
    if (k > 0 && seg.show_one(k - 1) < infl) {
      ans += seg.show_one(k - 1) + x + d;

      seg.update(idx[k - 1][0],
                 min(x + (k - idx[k - 1][0]) * d, seg.show_one(idx[k - 1][0])));
      if (idx[k - 1][0] < k - 1) {
        seg.update(k - 1, infl);
      }
      idx[idx[k][1]][0] = idx[k - 1][0];
      idx[k][0] = idx[k - 1][0];
      idx[k - 1][1] = idx[k][1];
      idx[idx[k - 1][0]][1] = idx[k][1];
      cnt--;
    }
    if (k < n && seg.show_one(k + 1) < infl) {
      ans += seg.show_one(k + 1) + x + d;
      seg.update(idx[k + 1][1],
                 min(x + (idx[k + 1][1] - k) * d, seg.show_one(idx[k + 1][1])));
      if (idx[k + 1][1] > k + 1) {
        seg.update(k + 1, infl);
      }
      idx[idx[k][0]][1] = idx[k + 1][1];
      idx[k][1] = idx[k + 1][1];
      idx[idx[k + 1][1]][0] = idx[k][0];
      idx[k + 1][0] = idx[k][0];
      cnt--;
    }
    seg.update(k, infl);
  }
  // rep(i, n) cout << seg.show_one(i) << " ";
  // cout << endl;
  cout << ans << endl;
  // seg.view();
}


Fixed submission:
#include <algorithm>
#include <bits/stdc++.h>
#include <bitset>
#include <cassert>
#include <iomanip>
#include <iostream>
#include <map>
#include <math.h>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <vector>
using namespace std;
#define rep(i, n) for (ll i = 0; i < (ll)n; i++)
#define repf(i, a, b) for (ll i = (ll)a; i < (ll)b; i++)
#define repr(i, a, b) for (ll i = (ll)a; i > (ll)b; i--)
#define repv(x, arr) for (auto &x : arr)
#define all(v) (v).begin(), (v).end()
#define vec(name, num) vector<ll> name((num), 0);
#define mp(a, b) make_pair(a, b)
#define pb(x) push_back(x)
#define F first
#define S second
typedef long long ll;
typedef long double ld;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;
typedef pair<ld, ld> pdd;
typedef vector<ll> vll;
typedef vector<vll> vvll;
typedef vector<int> vii;
typedef vector<vii> vvii;
const ll mod = 1e9 + 7;
const int infi = 2147483600;
const ll infl = 1e17;

struct SegmentTree {
private:
  int N;
  vector<ll> node;
  ll e;

public:
  SegmentTree() {}
  SegmentTree(ll e) : e(e) {}
  void __init__(int sz) {
    N = 1;
    while (N < sz)
      N *= 2;
    node.resize(2 * N - 1, e);
  }
  void build(vector<ll> &v) {
    int sz = (int)(v.size());
    __init__(sz);
    rep(i, sz) { node[i + N - 1] = v[i]; }
    for (int i = N - 2; i >= 0; i--) {
      node[i] = min(node[i * 2 + 1], node[i * 2 + 2]);
    }
  }
  void update(int k, ll x) {
    k += N - 1;
    node[k] = x;
    while (k > 0) {
      k = (k - 1) / 2;
      node[k] = min(node[2 * k + 1], node[2 * k + 2]);
    }
    return;
  }
  ll query(int a, int b) { return query(a, b, 0, 0, N); }
  ll query(int a, int b, int k, int l, int r) {
    if (b <= l || r <= a)
      return e;
    if (a <= l && r <= b)
      return node[k];
    ll vl, vr;
    vl = query(a, b, 2 * k + 1, l, (l + r) / 2);
    vr = query(a, b, 2 * k + 2, (l + r) / 2, r);
    return min(vl, vr);
  }
  ll show_one(int i) {
    if (0 <= i && i < N)
      return node[i + N - 1];
    return e;
  }

  int need() {
    int k = 0;
    while (k < N - 1) {
      node[k] = min(node[2 * k + 1], node[2 * k + 2]);
      if (node[k] == node[2 * k + 1]) {
        k = 2 * k + 1;
      } else if (node[k] == node[2 * k + 2]) {
        k = 2 * k + 2;
      }
    }
    return k - N + 1;
  }
  void view() {
    repf(k, N - 1, 2 * N - 1) {
      if (node[k] == e)
        cout << "inf ";
      else
        cout << node[k] << " ";
    }
    cout << endl;
    return;
  }
};

struct UnionFind {
  vector<int> root;
  int n;
  void __init__(int sz) {
    n = sz;
    root.resize(n, -1);
  }
  int find(int x) {
    if (root[x] >= 0) {
      root[x] = find(root[x]);
      return root[x];
    }
    return x;
  }
  void unit(int a, int b) {
    int ra, rb;
    ra = find(a);
    rb = find(b);
    if (ra == rb) {
      return;
    } else if (root[ra] <= root[rb]) {
      root[ra] += root[rb];
      root[rb] = ra;
    } else {
      root[rb] += root[ra];
      root[ra] = rb;
    }
    return;
  }
  int count() {
    set<int> cnt;
    rep(i, n) cnt.insert(find(i));
    return cnt.size();
  }
};

int main() {
  ll n, d;
  cin >> n >> d;
  vll a(n);
  rep(i, n) cin >> a[i];
  SegmentTree seg(infl);
  seg.build(a);
  ll ans = 0;
  ll cnt = n - 1;
  vvll idx(n, vll(2));
  rep(i, n) {
    idx[i][0] = i;
    idx[i][1] = i;
  }
  while (cnt > 0) {
    /*cout << ans;
    cout << endl;
    seg.view(); */
    int k = seg.need();
    ll x = seg.show_one(k);
    // cout << k << " " << x << " ";
    if (k > 0 && seg.show_one(k - 1) < infl) {
      ans += seg.show_one(k - 1) + x + d;

      seg.update(idx[k - 1][0],
                 min(x + (k - idx[k - 1][0]) * d, seg.show_one(idx[k - 1][0])));
      if (idx[k - 1][0] < k - 1) {
        seg.update(k - 1, infl);
      }
      idx[idx[k][1]][0] = idx[k - 1][0];
      idx[k][0] = idx[k - 1][0];
      idx[k - 1][1] = idx[k][1];
      idx[idx[k - 1][0]][1] = idx[k][1];
      cnt--;
    }
    if (k < n && seg.show_one(k + 1) < infl) {
      ans += seg.show_one(k + 1) + x + d;
      seg.update(idx[k + 1][1],
                 min(x + (idx[k + 1][1] - k) * d, seg.show_one(idx[k + 1][1])));
      if (idx[k + 1][1] > k + 1) {
        seg.update(k + 1, infl);
      }
      idx[idx[k][0]][1] = idx[k + 1][1];
      idx[k][1] = idx[k + 1][1];
      idx[idx[k + 1][1]][0] = idx[k][0];
      idx[k + 1][0] = idx[k][0];
      cnt--;
    }
    seg.update(k, infl);
  }
  // rep(i, n) cout << seg.show_one(i) << " ";
  // cout << endl;
  cout << ans << endl;
  // seg.view();
}