Problem p03153 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <algorithm>
#include <bitset>
#include <cassert>
#include <chrono>
#include <cmath>
#include <deque>
#include <iomanip>
#include <iostream>
#include <list>
#include <map>
#include <math.h>
#include <queue>
#include <random>
#include <set>
#include <sstream>
#include <string>
#include <time.h>
#include <unordered_map>
#include <unordered_set>
#include <vector>

using namespace std;

typedef long long ll;

#ifdef ONPC
mt19937 rnd(228);
#else
mt19937 rnd(chrono::high_resolution_clock::now().time_since_epoch().count());
#endif

const int N = 2e5 + 7;

int par[N];

int get(int v) {
  if (v == par[v]) {
    return v;
  } else {
    return par[v] = get(par[v]);
  }
}

void uni(int u, int v) { par[get(u)] = get(v); }

const ll inf = 1e18 + 7;

int main() {
#ifdef ONPC
  freopen("a.in", "r", stdin);
#endif
  ios::sync_with_stdio(0);
  cin.tie(0);
  int n, d;
  cin >> n >> d;
  vector<int> a(n);
  vector<int> ind(n);
  for (int i = 0; i < n; i++) {
    cin >> a[i];
    ind[i] = i;
  }
  for (int i = 0; i < n; i++) {
    par[i] = i;
  }
  int comp = n;
  ll ans = 0;
  while (comp > 1) {
    vector<ll> best(n, inf);
    vector<int> him(n, -1);
    auto rlx = [&]() {
      vector<pair<ll, int>> arr;
      for (int i = 0; i < n; i++) {
        int real_ind = ind[i];
        ll value = -i * (ll)d + a[real_ind];
        int who = get(real_ind);
        arr.push_back({value, who});
        sort(arr.begin(), arr.end());
        vector<pair<ll, int>> now;
        for (int i = 0; i < (int)arr.size(); i++) {
          if (now.empty() || now.back().second != arr[i].second) {
            now.push_back(arr[i]);
          }
        }
        while (now.size() > 2)
          now.pop_back();
        for (auto c : now) {
          if (c.second != who) {
            ll relax_cur = i * (ll)d + a[real_ind] + c.first;
            int relax_who = c.second;
            if (relax_cur < best[who]) {
              best[who] = relax_cur;
              him[who] = relax_who;
            }
          }
        }
      }
    };
    rlx();
    reverse(ind.begin(), ind.end());
    rlx();
    for (int i = 0; i < n; i++) {
      if (get(i) == i) {
        if (get(him[i]) != i) {
          uni(him[i], i);
          ans += best[i];
          comp--;
        }
      }
    }
  }
  cout << ans << '\n';
}


Fixed submission:
#include <algorithm>
#include <bitset>
#include <cassert>
#include <chrono>
#include <cmath>
#include <deque>
#include <iomanip>
#include <iostream>
#include <list>
#include <map>
#include <math.h>
#include <queue>
#include <random>
#include <set>
#include <sstream>
#include <string>
#include <time.h>
#include <unordered_map>
#include <unordered_set>
#include <vector>

using namespace std;

typedef long long ll;

#ifdef ONPC
mt19937 rnd(228);
#else
mt19937 rnd(chrono::high_resolution_clock::now().time_since_epoch().count());
#endif

const int N = 2e5 + 7;

int par[N];

int get(int v) {
  if (v == par[v]) {
    return v;
  } else {
    return par[v] = get(par[v]);
  }
}

void uni(int u, int v) { par[get(u)] = get(v); }

const ll inf = 1e18 + 7;

int main() {
#ifdef ONPC
  freopen("a.in", "r", stdin);
#endif
  ios::sync_with_stdio(0);
  cin.tie(0);
  int n, d;
  cin >> n >> d;
  vector<int> a(n);
  vector<int> ind(n);
  for (int i = 0; i < n; i++) {
    cin >> a[i];
    ind[i] = i;
  }
  for (int i = 0; i < n; i++) {
    par[i] = i;
  }
  int comp = n;
  ll ans = 0;
  while (comp > 1) {
    vector<ll> best(n, inf);
    vector<int> him(n, -1);
    auto rlx = [&]() {
      vector<pair<ll, int>> arr;
      for (int i = 0; i < n; i++) {
        int real_ind = ind[i];
        ll value = -i * (ll)d + a[real_ind];
        int who = get(real_ind);
        arr.push_back({value, who});
        sort(arr.begin(), arr.end());
        vector<pair<ll, int>> now;
        for (int i = 0; i < (int)arr.size(); i++) {
          if (now.empty() || now.back().second != arr[i].second) {
            now.push_back(arr[i]);
          }
        }
        while (now.size() > 2)
          now.pop_back();
        arr = now;
        for (auto c : now) {
          if (c.second != who) {
            ll relax_cur = i * (ll)d + a[real_ind] + c.first;
            int relax_who = c.second;
            if (relax_cur < best[who]) {
              best[who] = relax_cur;
              him[who] = relax_who;
            }
          }
        }
      }
    };
    rlx();
    reverse(ind.begin(), ind.end());
    rlx();
    for (int i = 0; i < n; i++) {
      if (get(i) == i) {
        if (get(him[i]) != i) {
          uni(him[i], i);
          ans += best[i];
          comp--;
        }
      }
    }
  }
  cout << ans << '\n';
}