Problem p03153 (C++)

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

// 標準ライブラリ関数の名前を変える
#define y0 y3487465
#define y1 y8687969
#define j0 j1347829
#define j1 j234892
#define next asdnext
#define prev asdprev

using ll = long long;
using ld = long double;
using pll = pair<ll, ll>;
using pld = pair<ld, ld>;
const int INF = 1e9 + 7;
const ll LINF = 1LL << 60;
const ll MOD = 1e9 + 7;
const ld PI = acos(-1);
const ld EPS = 1e-9; // 微調整用（EPSより小さいと0と判定など）

#define gcd __gcd // llは受け取ってくれない
int lcm(int a, int b) { return a / gcd(a, b) * b; }
#define ALL(a) a.begin(), a.end() // sort(ALL(vec));
#define REP(i, m, n) for (int i = (int)(m); i < (int)(n); ++i)
#define rep(i, n) REP(i, 0, n)
#define MP(a, b) make_pair(a, b)
#define SORT_UNIQUE(c)                                                         \
  (sort(c.begin(), c.end()),                                                   \
   c.resize(distance(c.begin(), unique(c.begin(), c.end()))))
#define GET_POS(c, x) (lower_bound(c.begin(), c.end(), x) - c.begin())
#define PB push_back
#define SZ(x) ((int)(x).size) // size()がunsignedなのでエラー避けに

// 最大値、最小値を更新する。aよりbのが大きい（小さい）か等しければaを更新してtrueを返す。そうでなければ何もせずfalseを返す
// chmax(nowmax,x);
template <typename T> bool chmax(T &a, T b) { return (a = max(a, b)) == b; }
template <typename T> bool chmin(T &a, T b) { return (a = min(a, b)) == b; }

// ------- template ------ //

// Range Minimum(Maximum) Query (RMQ)
// Big Range Query
// セグメント木 pll(x, y) = pll(値, 要素番号)
// インデックス取得したい場合: get_m**_with_idx()を使う
class SegmentTree {
public:
  int n = 1;
  vector<pll> seg[2]; // 0:最大	1:最小
  void init(int N) {  // N要素のセグメントツリー作成
    while (n < N)
      n *= 2;
    rep(i, 2) seg[i].resize(2 * n);
    fill(ALL(seg[0]), pll(-LINF, -1));
    fill(ALL(seg[1]), pll(LINF, -1));
  }
  pll cmp_min(pll x, pll y) { // 比較関数
    if (x.first == y.first)
      return x.second < y.second ? x : y;
    return min(x, y);
  }
  pll cmp_max(pll x, pll y) { // 比較関数
    if (x.first == y.first)
      return x.second < y.second ? x : y;
    return max(x, y);
  }
  void update(ll k, ll x, bool flag = false) { // idx:k番目の値をxに変更
    k += n - 1;
    seg[0][k] = seg[1][k] = pll(x, k - (n - 1));
    if (flag) {
      seg[0][k] = pll(-LINF, -1);
      seg[1][k] = pll(LINF, -1);
    }
    while (k > 0) {
      k = (k - 1) / 2;
      seg[0][k] = cmp_max(seg[0][k * 2 + 1], seg[0][k * 2 + 2]);
      seg[1][k] = cmp_min(seg[1][k * 2 + 1], seg[1][k * 2 + 2]);
    }
  }
  pll calc(ll a, ll b, ll k, ll l, ll r, ll f) {
    if (b <= l || r <= a)
      return f ? pll(LINF, -1) : pll(-LINF, -1);
    if (a <= l && r <= b)
      return seg[f][k];
    pll vl = calc(a, b, k * 2 + 1, l, (l + r) / 2, f);
    pll vr = calc(a, b, k * 2 + 2, (l + r) / 2, r, f);
    return f ? cmp_min(vl, vr) : cmp_max(vl, vr);
  }
  void erase(ll k) {
    update(k, -1, true);
  } // 要素を削除（実際にはINF, -INFにする）
  // 最大値・最小値のみ取得
  ll get_max(ll a, ll b) { return calc(a, b, 0, 0, n, 0).first; }
  ll get_min(ll a, ll b) { return calc(a, b, 0, 0, n, 1).first; }
  // インデックスも取得（pll<最大値, 最大値のidx>）を返す
  pll get_max_with_idx(ll a, ll b) { return calc(a, b, 0, 0, n, 0); }
  pll get_min_with_idx(ll a, ll b) { return calc(a, b, 0, 0, n, 1); }
};

// Union-Find（素集合データ構造）
struct UnionFind {
  vector<int> data;

  UnionFind(int sz) { data.assign(sz, -1); }

  bool unite(int x, int y) {
    x = find(x), y = find(y);
    if (x == y)
      return (false);
    if (data[x] > data[y])
      swap(x, y);
    data[x] += data[y];
    data[y] = x;
    return (true);
  }

  int find(int k) {
    if (data[k] < 0)
      return (k);
    return (data[k] = find(data[k]));
  }

  int size(int k) { return (-data[find(k)]); }
};

int main() {
  ll N, D;
  cin >> N >> D;

  // 1. Asに[a:都市の規模, i:idx]を格納
  vector<pll> As(N); // (ll 規模, ll idx)
  rep(i, N) {
    cin >> As.at(i).first;
    As.at(i).second = i;
  }

  // 2. 規模降順にソート
  sort(ALL(As), greater<pll>()); // 降順

  /*  check
    for (pll a: As){
      cout << a.first << ' ' << a.second << endl;
    }   */

  // 3.
  // SegTreeを2つ作り、それぞれにコストを入れる（segtreeのidxがAsにおけるidxとなるようにする）
  SegmentTree seg_left;
  seg_left.init(N);
  SegmentTree seg_right;
  seg_right.init(N);

  int i = 0;
  for (pll a : As) {
    seg_left.update(a.second, (N - 1 - a.second) * D + a.first);
    seg_right.update(a.second, a.second * D + a.first);
    i++;
  }
  // cout << seg_left.get_min_with_idx(1,3).first << ' ' <<
  // seg_left.get_min_with_idx(1,3).second << endl; cout <<
  // seg_left.get_max_with_idx(2,6).first << ' ' <<
  // seg_left.get_max_with_idx(2,6).second << endl;

  // ここからの流れ
  // 4. それを使って、規模降順forで左右領域の「[l,
  // r)における最もコストの低い接続先都市」をピックアップする
  //   --> ピックアップしたら、priority_queue(edge_candidates)に[idx,
  //   adjidx]を入れていく

  // edge_candidatesを宣言 最小値を取り出すpriority_queue
  priority_queue<tuple<ll, ll, ll>, vector<tuple<ll, ll, ll>>,
                 greater<tuple<ll, ll, ll>>>
      edge_candidates;
  ll nowidx, nowcost;
  ll leftmin, rightmin, lminidx, rminidx;
  for (pll a : As) {
    tie(nowcost, nowidx) = a;
    //    cout << nowcost << ' ' << nowidx << ' ';
    if (nowidx == 0) {
      tie(rightmin, rminidx) = seg_right.get_min_with_idx(1, N);
      if (rightmin !=
          INF) { // INFなら、基点となる都市より規模の大きい都市なので、無視
        edge_candidates.push(make_tuple(
            abs(nowidx - rminidx) * D + nowcost + rightmin - rminidx * D,
            nowidx, rminidx)); //(コスト, a_idx, adj_idx)
      }
    } else if (nowidx == N - 1) {
      tie(leftmin, lminidx) = seg_left.get_min_with_idx(0, nowidx);
      if (leftmin != LINF) {
        edge_candidates.push(make_tuple(abs(nowidx - lminidx) * D + nowcost +
                                            leftmin - (N - 1 - lminidx) * D,
                                        nowidx,
                                        lminidx)); //(コスト, a_idx, adj_idx)
      }
    } else {
      tie(leftmin, lminidx) = seg_left.get_min_with_idx(0, nowidx);
      if (leftmin != LINF) {
        edge_candidates.push(make_tuple(abs(nowidx - lminidx) * D + nowcost +
                                            leftmin - (N - 1 - lminidx) * D,
                                        nowidx,
                                        lminidx)); //(コスト, a_idx, adj_idx)
      }
      tie(rightmin, rminidx) = seg_right.get_min_with_idx(nowidx + 1, N);
      if (rightmin !=
          LINF) { // INFなら、基点となる都市より規模の大きい都市なので、無視
        edge_candidates.push(make_tuple(
            abs(nowidx - rminidx) * D + nowcost + rightmin - rminidx * D,
            nowidx, rminidx)); //(コスト, a_idx, adj_idx)
      }
    }
    //   cout << leftmin << ' ' << lminidx << ' ';
    //   cout << rightmin << ' ' << rminidx << endl;
    seg_left.erase(nowidx);
    seg_right.erase(nowidx);
  }
  // 5. 1つずつtop()->pop()して、Uniteする（Uniteできたらans+=コスト,
  // できなかったら無視）

  // check
  /*ll aa, bb, cc;
   tie(aa, bb, cc) = edge_candidates.top();
   cout << "1ban : " << aa << ' ' << bb << ' ' << cc << endl;
   edge_candidates.pop();
   tie(aa, bb, cc) = edge_candidates.top();
   cout << "2ban : " << aa << ' ' << bb << ' ' << cc << endl; */

  UnionFind UF(N);
  ll ans = 0;
  ll cost, sidx, adjidx;
  while (not edge_candidates.empty()) {
    //    cout << 'a' << endl;
    tie(cost, sidx, adjidx) = edge_candidates.top();
    //    cout << cost << ' ' << sidx << ' ' << adjidx << endl;
    edge_candidates.pop();
    if (UF.unite(sidx, adjidx)) { // uniteが成功したら=まだ接続されていなかったら
      ans += cost;
    } // 接続されていたら無視する
  }

  //  cout << "end" << endl;
  // 6. できあがり！
  cout << ans << endl;
}


Fixed submission:
#include <bits/stdc++.h>
using namespace std;

// 標準ライブラリ関数の名前を変える
#define y0 y3487465
#define y1 y8687969
#define j0 j1347829
#define j1 j234892
#define next asdnext
#define prev asdprev

using ll = long long;
using ld = long double;
using pll = pair<ll, ll>;
using pld = pair<ld, ld>;
const int INF = 1e9 + 7;
const ll LINF = 1LL << 60;
const ll MOD = 1e9 + 7;
const ld PI = acos(-1);
const ld EPS = 1e-9; // 微調整用（EPSより小さいと0と判定など）

#define gcd __gcd // llは受け取ってくれない
int lcm(int a, int b) { return a / gcd(a, b) * b; }
#define ALL(a) a.begin(), a.end() // sort(ALL(vec));
#define REP(i, m, n) for (int i = (int)(m); i < (int)(n); ++i)
#define rep(i, n) REP(i, 0, n)
#define MP(a, b) make_pair(a, b)
#define SORT_UNIQUE(c)                                                         \
  (sort(c.begin(), c.end()),                                                   \
   c.resize(distance(c.begin(), unique(c.begin(), c.end()))))
#define GET_POS(c, x) (lower_bound(c.begin(), c.end(), x) - c.begin())
#define PB push_back
#define SZ(x) ((int)(x).size) // size()がunsignedなのでエラー避けに

// 最大値、最小値を更新する。aよりbのが大きい（小さい）か等しければaを更新してtrueを返す。そうでなければ何もせずfalseを返す
// chmax(nowmax,x);
template <typename T> bool chmax(T &a, T b) { return (a = max(a, b)) == b; }
template <typename T> bool chmin(T &a, T b) { return (a = min(a, b)) == b; }

// ------- template ------ //

// Range Minimum(Maximum) Query (RMQ)
// Big Range Query
// セグメント木 pll(x, y) = pll(値, 要素番号)
// インデックス取得したい場合: get_m**_with_idx()を使う
class SegmentTree {
public:
  int n = 1;
  vector<pll> seg[2]; // 0:最大	1:最小
  void init(int N) {  // N要素のセグメントツリー作成
    while (n < N)
      n *= 2;
    rep(i, 2) seg[i].resize(2 * n);
    fill(ALL(seg[0]), pll(-LINF, -1));
    fill(ALL(seg[1]), pll(LINF, -1));
  }
  pll cmp_min(pll x, pll y) { // 比較関数
    if (x.first == y.first)
      return x.second < y.second ? x : y;
    return min(x, y);
  }
  pll cmp_max(pll x, pll y) { // 比較関数
    if (x.first == y.first)
      return x.second < y.second ? x : y;
    return max(x, y);
  }
  void update(ll k, ll x, bool flag = false) { // idx:k番目の値をxに変更
    k += n - 1;
    seg[0][k] = seg[1][k] = pll(x, k - (n - 1));
    if (flag) {
      seg[0][k] = pll(-LINF, -1);
      seg[1][k] = pll(LINF, -1);
    }
    while (k > 0) {
      k = (k - 1) / 2;
      seg[0][k] = cmp_max(seg[0][k * 2 + 1], seg[0][k * 2 + 2]);
      seg[1][k] = cmp_min(seg[1][k * 2 + 1], seg[1][k * 2 + 2]);
    }
  }
  pll calc(ll a, ll b, ll k, ll l, ll r, ll f) {
    if (b <= l || r <= a)
      return f ? pll(LINF, -1) : pll(-LINF, -1);
    if (a <= l && r <= b)
      return seg[f][k];
    pll vl = calc(a, b, k * 2 + 1, l, (l + r) / 2, f);
    pll vr = calc(a, b, k * 2 + 2, (l + r) / 2, r, f);
    return f ? cmp_min(vl, vr) : cmp_max(vl, vr);
  }
  void erase(ll k) {
    update(k, -1, true);
  } // 要素を削除（実際にはINF, -INFにする）
  // 最大値・最小値のみ取得
  ll get_max(ll a, ll b) { return calc(a, b, 0, 0, n, 0).first; }
  ll get_min(ll a, ll b) { return calc(a, b, 0, 0, n, 1).first; }
  // インデックスも取得（pll<最大値, 最大値のidx>）を返す
  pll get_max_with_idx(ll a, ll b) { return calc(a, b, 0, 0, n, 0); }
  pll get_min_with_idx(ll a, ll b) { return calc(a, b, 0, 0, n, 1); }
};

// Union-Find（素集合データ構造）
struct UnionFind {
  vector<int> data;

  UnionFind(int sz) { data.assign(sz, -1); }

  bool unite(int x, int y) {
    x = find(x), y = find(y);
    if (x == y)
      return (false);
    if (data[x] > data[y])
      swap(x, y);
    data[x] += data[y];
    data[y] = x;
    return (true);
  }

  int find(int k) {
    if (data[k] < 0)
      return (k);
    return (data[k] = find(data[k]));
  }

  int size(int k) { return (-data[find(k)]); }
};

int main() {
  ll N, D;
  cin >> N >> D;

  // 1. Asに[a:都市の規模, i:idx]を格納
  vector<pll> As(N); // (ll 規模, ll idx)
  rep(i, N) {
    cin >> As.at(i).first;
    As.at(i).second = i;
  }

  // 2. 規模降順にソート
  sort(ALL(As), greater<pll>()); // 降順

  /*  check
    for (pll a: As){
      cout << a.first << ' ' << a.second << endl;
    }   */

  // 3.
  // SegTreeを2つ作り、それぞれにコストを入れる（segtreeのidxがAsにおけるidxとなるようにする）
  SegmentTree seg_left;
  seg_left.init(N);
  SegmentTree seg_right;
  seg_right.init(N);

  int i = 0;
  for (pll a : As) {
    seg_left.update(a.second, (N - 1 - a.second) * D + a.first);
    seg_right.update(a.second, a.second * D + a.first);
    i++;
  }
  // cout << seg_left.get_min_with_idx(1,3).first << ' ' <<
  // seg_left.get_min_with_idx(1,3).second << endl; cout <<
  // seg_left.get_max_with_idx(2,6).first << ' ' <<
  // seg_left.get_max_with_idx(2,6).second << endl;

  // ここからの流れ
  // 4. それを使って、規模降順forで左右領域の「[l,
  // r)における最もコストの低い接続先都市」をピックアップする
  //   --> ピックアップしたら、priority_queue(edge_candidates)に[idx,
  //   adjidx]を入れていく

  // edge_candidatesを宣言 最小値を取り出すpriority_queue
  priority_queue<tuple<ll, ll, ll>, vector<tuple<ll, ll, ll>>,
                 greater<tuple<ll, ll, ll>>>
      edge_candidates;
  ll nowidx, nowcost;
  ll leftmin, rightmin, lminidx, rminidx;
  for (pll a : As) {
    tie(nowcost, nowidx) = a;
    //    cout << nowcost << ' ' << nowidx << ' ';
    if (nowidx == 0) {
      tie(rightmin, rminidx) = seg_right.get_min_with_idx(1, N);
      if (rightmin !=
          LINF) { // INFなら、基点となる都市より規模の大きい都市なので、無視
        edge_candidates.push(make_tuple(
            abs(nowidx - rminidx) * D + nowcost + rightmin - rminidx * D,
            nowidx, rminidx)); //(コスト, a_idx, adj_idx)
      }
    } else if (nowidx == N - 1) {
      tie(leftmin, lminidx) = seg_left.get_min_with_idx(0, nowidx);
      if (leftmin != LINF) {
        edge_candidates.push(make_tuple(abs(nowidx - lminidx) * D + nowcost +
                                            leftmin - (N - 1 - lminidx) * D,
                                        nowidx,
                                        lminidx)); //(コスト, a_idx, adj_idx)
      }
    } else {
      tie(leftmin, lminidx) = seg_left.get_min_with_idx(0, nowidx);
      if (leftmin != LINF) {
        edge_candidates.push(make_tuple(abs(nowidx - lminidx) * D + nowcost +
                                            leftmin - (N - 1 - lminidx) * D,
                                        nowidx,
                                        lminidx)); //(コスト, a_idx, adj_idx)
      }
      tie(rightmin, rminidx) = seg_right.get_min_with_idx(nowidx + 1, N);
      if (rightmin !=
          LINF) { // INFなら、基点となる都市より規模の大きい都市なので、無視
        edge_candidates.push(make_tuple(
            abs(nowidx - rminidx) * D + nowcost + rightmin - rminidx * D,
            nowidx, rminidx)); //(コスト, a_idx, adj_idx)
      }
    }
    //   cout << leftmin << ' ' << lminidx << ' ';
    //   cout << rightmin << ' ' << rminidx << endl;
    seg_left.erase(nowidx);
    seg_right.erase(nowidx);
  }
  // 5. 1つずつtop()->pop()して、Uniteする（Uniteできたらans+=コスト,
  // できなかったら無視）

  // check
  /*ll aa, bb, cc;
   tie(aa, bb, cc) = edge_candidates.top();
   cout << "1ban : " << aa << ' ' << bb << ' ' << cc << endl;
   edge_candidates.pop();
   tie(aa, bb, cc) = edge_candidates.top();
   cout << "2ban : " << aa << ' ' << bb << ' ' << cc << endl; */

  UnionFind UF(N);
  ll ans = 0;
  ll cost, sidx, adjidx;
  while (not edge_candidates.empty()) {
    //    cout << 'a' << endl;
    tie(cost, sidx, adjidx) = edge_candidates.top();
    //    cout << cost << ' ' << sidx << ' ' << adjidx << endl;
    edge_candidates.pop();
    if (UF.unite(sidx, adjidx)) { // uniteが成功したら=まだ接続されていなかったら
      ans += cost;
    } // 接続されていたら無視する
  }

  //  cout << "end" << endl;
  // 6. できあがり！
  cout << ans << endl;
}